<commit_before>package main

// This file contains about everything related to posts. At the top you will find routes
// and at the bottom you can find CRUD options. The functions in this file are analogous
// to the ones in users.go, although some differences exist.

import (
	"bufio"
	"bytes"
	"errors"
	r "github.com/dancannon/gorethink"
	"github.com/go-martini/martini"
	"github.com/gosimple/slug"
	"github.com/martini-contrib/render"
	"github.com/martini-contrib/sessions"
	"log"
	"net/http"
	"strings"
	"time"
)

type Post struct {
	Date      int32  `json:"date" gorethink:"date"`
	Title     string `json:"title"form:"title" binding:"required" gorethink:"title"`
	Author    string `json:"author,omitempty" gorethink:"author"`
	Content   string `json:",omitempty" form:"content" binding:"required" gorethink:"content"`
	Excerpt   string `json:"excerpt" gorethink:"excerpt"`
	Slug      string `json:"slug" gorethink:"slug"`
	Published bool   `json:"-" gorethink:"published"`
}

// Generates 15 word excerpt from given input.
func Excerpt(input string) string {
	scanner := bufio.NewScanner(strings.NewReader(input))
	scanner.Split(bufio.ScanWords)
	count := 0
	var excerpt bytes.Buffer
	for scanner.Scan() && count < 15 {
		count++
		excerpt.WriteString(scanner.Text() + " ")
	}
	return strings.TrimSpace(excerpt.String())
}

func CreatePost(req *http.Request, s sessions.Session, db *r.Session, res render.Render, post Post) {
	entry, err := post.Insert(db, s)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		return
	}
	switch root(req) {
	case "api":
		res.JSON(200, entry)
		return
	case "post":
		res.Redirect("/user", 302)
		return
	}
	res.JSON(500, map[string]interface{}{"error": "Internal server error"})
}

func ReadPosts(res render.Render, db *r.Session) {
	var post Post
	posts, err := post.GetAll(db)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		return
	}
	res.JSON(200, posts)
}

func ReadPost(req *http.Request, params martini.Params, res render.Render, db *r.Session) {
	var post Post
	if params["title"] == "new" {
		res.JSON(406, map[string]interface{}{"error": "You cant name a post with colliding route name!"})
		return
	}
	post.Slug = params["title"]
	post, err := post.Get(db)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	switch root(req) {
	case "api":
		res.JSON(200, post)
		return
	case "post":
		res.HTML(200, "post/display", post)
		return
	}
	res.JSON(500, map[string]interface{}{"error": "Internal server error"})
}

func EditPost(req *http.Request, params martini.Params, res render.Render, db *r.Session) {
	var post Post
	if params["title"] == "new" {
		res.JSON(406, map[string]interface{}{"error": "You cant name a post with colliding route name!"})
		return
	}
	post.Slug = params["title"]
	post, err := post.Get(db)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	switch root(req) {
	case "api":
		res.JSON(200, post)
		return
	case "post":
		res.HTML(200, "post/edit", post)
		return
	}
	res.JSON(500, map[string]interface{}{"error": "Internal server error"})
}

func UpdatePost(req *http.Request, params martini.Params, s sessions.Session, res render.Render, db *r.Session, post Post) {
	post.Slug = params["title"]
	entry, err := post.Get(db)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	post, err = entry.Update(db, s, post)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	switch root(req) {
	case "api":
		res.JSON(200, post)
		return
	case "post":
		res.Redirect("/user", 302)
		return
	}
	res.JSON(500, map[string]interface{}{"error": "Internal server error"})
}

func PublishPost(req *http.Request, params martini.Params, s sessions.Session, res render.Render, db *r.Session) {
	var post Post
	post.Slug = params["title"]
	post, err := post.Get(db)
	post.Published = true
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	post, err = post.Update(db, s, post)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	switch root(req) {
	case "api":
		res.JSON(200, map[string]interface{}{"success": "Post published"})
		return
	case "post":
		res.Redirect("/post/"+post.Slug, 302)
		return
	}
	res.JSON(500, map[string]interface{}{"error": "Internal server error"})
}

func DeletePost(req *http.Request, params martini.Params, s sessions.Session, res render.Render, db *r.Session) {
	var post Post
	post.Slug = params["title"]
	post, err := post.Get(db)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	err = post.Delete(db, s)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	switch root(req) {
	case "api":
		res.JSON(200, map[string]interface{}{"success": "Post deleted"})
		return
	case "post":
		res.Redirect("/user", 302)
		return
	}
	res.JSON(500, map[string]interface{}{"error": "Internal server error"})
}

// Insert inserts Post object into database. The function fetches user data from session and
// after that it fills Author, Date and Excerpt fields. Returns created Post object.
func (post Post) Insert(db *r.Session, s sessions.Session) (Post, error) {
	var person Person
	person, err := person.Session(db, s)
	if err != nil {
		return post, err
	}
	post.Author = person.Id
	post.Date = int32(time.Now().Unix())
	post.Excerpt = Excerpt(post.Content)
	post.Slug = slug.Make(post.Title)
	post.Published = false
	post.Content = strings.TrimPrefix(post.Content, "

")
	post.Content = strings.TrimSuffix(post.Content, "

")
	post.Content = strings.Replace(post.Content, "

", "", -1)
	post.Content = strings.Replace(post.Content, "
", "", -1)
	row, err := r.Table("posts").Insert(post).RunRow(db)
	if err != nil {
		return post, err
	}
	err = row.Scan(&post)
	if err != nil {
		return post, err
	}
	return post, err
}

func (post Post) Get(s *r.Session) (Post, error) {
	row, err := r.Table("posts").Filter(func(this r.RqlTerm) r.RqlTerm {
		return this.Field("slug").Eq(post.Slug)
	}).RunRow(s)
	if err != nil {
		return post, err
	}
	if row.IsNil() {
		return post, errors.New("Nothing was found.")
	}
	err = row.Scan(&post)
	if err != nil {
		return post, err
	}
	return post, err
}

// Update updates a `entry` paramter with `post` parameter. Returns updated Post object.
func (entry Post) Update(db *r.Session, s sessions.Session, post Post) (Post, error) {
	var person Person
	person, err := person.Session(db, s)
	if err != nil {
		return post, err
	}
	if entry.Author == person.Id {
		row, err := r.Table("posts").Filter(func(this r.RqlTerm) r.RqlTerm {
			return this.Field("slug").Eq(entry.Slug)
		}).Update(map[string]interface{}{"published": post.Published, "content": post.Content, "slug": slug.Make(post.Title), "title": post.Title, "excerpt": Excerpt(post.Content)}).RunRow(db)
		if err != nil {
			return post, err
		}
		if row.IsNil() {
			return post, errors.New("Nothing was found.")
		}
		err = row.Scan(&post)
		if err != nil {
			return post, err
		}
	} else {
		return post, errors.New("Unauthorized.")
	}
	return post, err
}

// Delete deletes a post.
func (post Post) Delete(db *r.Session, s sessions.Session) error {
	var person Person
	person, err := person.Session(db, s)
	if err != nil {
		return err
	}
	if post.Author == person.Id {
		row, err := r.Table("posts").Filter(func(this r.RqlTerm) r.RqlTerm {
			return this.Field("slug").Eq(post.Slug)
		}).Delete().RunRow(db)
		if err != nil {
			return err
		}
		if row.IsNil() {
			return errors.New("Nothing was found.")
		}
		err = row.Scan(&post)
		if err != nil {
			return err
		}
	} else {
		return errors.New("Unauthorized.")
	}
	return err
}

func (post Post) GetAll(s *r.Session) ([]Post, error) {
	var posts []Post
	rows, err := r.Table("posts").OrderBy(r.Desc("date")).Run(s)
	if err != nil {
		return nil, err
	}
	for rows.Next() {
		err := rows.Scan(&post)
		post, err := post.Get(s)
		if err != nil {
			return nil, err
		}
		if post.Published {
			posts = append(posts, post)
		}
	}
	return posts, nil
}
<commit_msg>remove hacky content manipulations<commit_after>package main

// This file contains about everything related to posts. At the top you will find routes
// and at the bottom you can find CRUD options. The functions in this file are analogous
// to the ones in users.go, although some differences exist.

import (
	"bufio"
	"bytes"
	"errors"
	r "github.com/dancannon/gorethink"
	"github.com/go-martini/martini"
	"github.com/gosimple/slug"
	"github.com/martini-contrib/render"
	"github.com/martini-contrib/sessions"
	"log"
	"net/http"
	"strings"
	"time"
)

type Post struct {
	Date      int32  `json:"date" gorethink:"date"`
	Title     string `json:"title"form:"title" binding:"required" gorethink:"title"`
	Author    string `json:"author,omitempty" gorethink:"author"`
	Content   string `json:",omitempty" form:"content" binding:"required" gorethink:"content"`
	Excerpt   string `json:"excerpt" gorethink:"excerpt"`
	Slug      string `json:"slug" gorethink:"slug"`
	Published bool   `json:"-" gorethink:"published"`
}

// Generates 15 word excerpt from given input.
func Excerpt(input string) string {
	scanner := bufio.NewScanner(strings.NewReader(input))
	scanner.Split(bufio.ScanWords)
	count := 0
	var excerpt bytes.Buffer
	for scanner.Scan() && count < 15 {
		count++
		excerpt.WriteString(scanner.Text() + " ")
	}
	return strings.TrimSpace(excerpt.String())
}

func CreatePost(req *http.Request, s sessions.Session, db *r.Session, res render.Render, post Post) {
	entry, err := post.Insert(db, s)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		return
	}
	switch root(req) {
	case "api":
		res.JSON(200, entry)
		return
	case "post":
		res.Redirect("/user", 302)
		return
	}
	res.JSON(500, map[string]interface{}{"error": "Internal server error"})
}

func ReadPosts(res render.Render, db *r.Session) {
	var post Post
	posts, err := post.GetAll(db)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		return
	}
	res.JSON(200, posts)
}

func ReadPost(req *http.Request, params martini.Params, res render.Render, db *r.Session) {
	var post Post
	if params["title"] == "new" {
		res.JSON(406, map[string]interface{}{"error": "You cant name a post with colliding route name!"})
		return
	}
	post.Slug = params["title"]
	post, err := post.Get(db)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	switch root(req) {
	case "api":
		res.JSON(200, post)
		return
	case "post":
		res.HTML(200, "post/display", post)
		return
	}
	res.JSON(500, map[string]interface{}{"error": "Internal server error"})
}

func EditPost(req *http.Request, params martini.Params, res render.Render, db *r.Session) {
	var post Post
	if params["title"] == "new" {
		res.JSON(406, map[string]interface{}{"error": "You cant name a post with colliding route name!"})
		return
	}
	post.Slug = params["title"]
	post, err := post.Get(db)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	switch root(req) {
	case "api":
		res.JSON(200, post)
		return
	case "post":
		res.HTML(200, "post/edit", post)
		return
	}
	res.JSON(500, map[string]interface{}{"error": "Internal server error"})
}

func UpdatePost(req *http.Request, params martini.Params, s sessions.Session, res render.Render, db *r.Session, post Post) {
	post.Slug = params["title"]
	entry, err := post.Get(db)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	post, err = entry.Update(db, s, post)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	switch root(req) {
	case "api":
		res.JSON(200, post)
		return
	case "post":
		res.Redirect("/user", 302)
		return
	}
	res.JSON(500, map[string]interface{}{"error": "Internal server error"})
}

func PublishPost(req *http.Request, params martini.Params, s sessions.Session, res render.Render, db *r.Session) {
	var post Post
	post.Slug = params["title"]
	post, err := post.Get(db)
	post.Published = true
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	post, err = post.Update(db, s, post)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	switch root(req) {
	case "api":
		res.JSON(200, map[string]interface{}{"success": "Post published"})
		return
	case "post":
		res.Redirect("/post/"+post.Slug, 302)
		return
	}
	res.JSON(500, map[string]interface{}{"error": "Internal server error"})
}

func DeletePost(req *http.Request, params martini.Params, s sessions.Session, res render.Render, db *r.Session) {
	var post Post
	post.Slug = params["title"]
	post, err := post.Get(db)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	err = post.Delete(db, s)
	if err != nil {
		res.JSON(500, map[string]interface{}{"error": "Internal server error"})
		log.Println(err)
		return
	}
	switch root(req) {
	case "api":
		res.JSON(200, map[string]interface{}{"success": "Post deleted"})
		return
	case "post":
		res.Redirect("/user", 302)
		return
	}
	res.JSON(500, map[string]interface{}{"error": "Internal server error"})
}

// Insert inserts Post object into database. The function fetches user data from session and
// after that it fills Author, Date and Excerpt fields. Returns created Post object.
func (post Post) Insert(db *r.Session, s sessions.Session) (Post, error) {
	var person Person
	person, err := person.Session(db, s)
	if err != nil {
		return post, err
	}
	post.Author = person.Id
	post.Date = int32(time.Now().Unix())
	post.Excerpt = Excerpt(post.Content)
	post.Slug = slug.Make(post.Title)
	post.Published = false
	row, err := r.Table("posts").Insert(post).RunRow(db)
	if err != nil {
		return post, err
	}
	err = row.Scan(&post)
	if err != nil {
		return post, err
	}
	return post, err
}

func (post Post) Get(s *r.Session) (Post, error) {
	row, err := r.Table("posts").Filter(func(this r.RqlTerm) r.RqlTerm {
		return this.Field("slug").Eq(post.Slug)
	}).RunRow(s)
	if err != nil {
		return post, err
	}
	if row.IsNil() {
		return post, errors.New("Nothing was found.")
	}
	err = row.Scan(&post)
	if err != nil {
		return post, err
	}
	return post, err
}

// Update updates a `entry` paramter with `post` parameter. Returns updated Post object.
func (entry Post) Update(db *r.Session, s sessions.Session, post Post) (Post, error) {
	var person Person
	person, err := person.Session(db, s)
	if err != nil {
		return post, err
	}
	if entry.Author == person.Id {
		row, err := r.Table("posts").Filter(func(this r.RqlTerm) r.RqlTerm {
			return this.Field("slug").Eq(entry.Slug)
		}).Update(map[string]interface{}{"published": post.Published, "content": post.Content, "slug": slug.Make(post.Title), "title": post.Title, "excerpt": Excerpt(post.Content)}).RunRow(db)
		if err != nil {
			return post, err
		}
		if row.IsNil() {
			return post, errors.New("Nothing was found.")
		}
		err = row.Scan(&post)
		if err != nil {
			return post, err
		}
	} else {
		return post, errors.New("Unauthorized.")
	}
	return post, err
}

// Delete deletes a post.
func (post Post) Delete(db *r.Session, s sessions.Session) error {
	var person Person
	person, err := person.Session(db, s)
	if err != nil {
		return err
	}
	if post.Author == person.Id {
		row, err := r.Table("posts").Filter(func(this r.RqlTerm) r.RqlTerm {
			return this.Field("slug").Eq(post.Slug)
		}).Delete().RunRow(db)
		if err != nil {
			return err
		}
		if row.IsNil() {
			return errors.New("Nothing was found.")
		}
		err = row.Scan(&post)
		if err != nil {
			return err
		}
	} else {
		return errors.New("Unauthorized.")
	}
	return err
}

func (post Post) GetAll(s *r.Session) ([]Post, error) {
	var posts []Post
	rows, err := r.Table("posts").OrderBy(r.Desc("date")).Run(s)
	if err != nil {
		return nil, err
	}
	for rows.Next() {
		err := rows.Scan(&post)
		post, err := post.Get(s)
		if err != nil {
			return nil, err
		}
		if post.Published {
			posts = append(posts, post)
		}
	}
	return posts, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"net/http"
	"owl/common/types"
	"strconv"
)

type userGroup struct {
	types.UserGroup
	Users []*types.User `json:"users" gorm:"many2many:user_user_group;"`
}

func userGroupList(c *gin.Context) {
	q := c.DefaultQuery("q", "")
	page, _ := strconv.Atoi(c.DefaultQuery("page", "1"))
	pageSize, _ := strconv.Atoi(c.DefaultQuery("pageSize", DefaultPageSize))

	groups := []*userGroup{}
	db := mydb.Table("user_group")
	if len(q) > 0 {
		q = fmt.Sprintf("%%%s%%", q)
		db = db.Where("name like ?", q)
	}
	offset := (page - 1) * pageSize
	db.Offset(offset).Limit(pageSize).Find(&groups)
	for _, group := range groups {
		mydb.Model(&group).Association("Users").Find(&group.Users)
	}

	c.JSON(http.StatusOK, gin.H{"code": http.StatusOK, "groups": groups})
}

func userGroupCreate(c *gin.Context) {
	group := userGroup{}
	response := gin.H{}
	defer c.JSON(http.StatusOK, response)
	if err := c.BindJSON(&group); err != nil {
		response["code"] = http.StatusBadRequest
		response["message"] = err.Error()
		return
	}
	var cnt int
	if mydb.Table("user_group").Where("name=?", group.Name).Count(&cnt); cnt != 0 {
		response["code"] = http.StatusBadRequest
		response["message"] = "user group already exists"
		return
	}
	if err := mydb.Set("gorm:save_associations", false).Save(&group).Error; err != nil {
		response["code"] = http.StatusInternalServerError
		response["message"] = err.Error()
		return
	}
	mydb.Model(&group).Association("Users").Replace(group.Users).Find(&group.Users)

	response["code"] = http.StatusOK
	response["message"] = "group create successful"
	response["group"] = group
}

func userGroupUpdate(c *gin.Context) {
	groups := []*userGroup{}
	response := gin.H{}
	defer c.JSON(http.StatusOK, response)
	if err := c.BindJSON(&groups); err != nil {
		response["code"] = http.StatusBadRequest
		response["message"] = err.Error()
		return
	}
	var message string
	for _, group := range groups {

		if mydb.Table("user_group").Where("id=?", group.ID).First(&group).RecordNotFound() {
			response["code"] = http.StatusBadRequest
			message += "user group not found\n"
			continue
		}
		if err := mydb.Set("gorm:save_associations", false).Save(&group).Association("Users").Replace(group.Users).Error; err != nil {
			response["code"] = http.StatusInternalServerError
			message += err.Error()
			continue
		}
	}
	response["message"] = message
	response["code"] = http.StatusOK
	response["groups"] = groups
	response["message"] = "user group update successful"
}

func userGroupDelete(c *gin.Context) {
	id, _ := strconv.Atoi(c.Param("id"))
	response := gin.H{}
	defer c.JSON(http.StatusOK, response)
	group := userGroup{}
	if mydb.Table("user_group").Where("id = ?", id).Find(&group).RecordNotFound() {
		response["code"] = http.StatusNotFound
		response["message"] = "The usergroup not found"
		return
	}
	if err := mydb.Delete(&group).Error; err != nil {
		response["code"] = http.StatusInternalServerError
		response["message"] = err.Error()
		return

	}
	response["code"] = http.StatusOK
	response["message"] = fmt.Sprintf("%s delete successful", group.Name)
}
<commit_msg>取消用户组分页<commit_after>package main

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"net/http"
	"owl/common/types"
	"strconv"
)

type userGroup struct {
	types.UserGroup
	Users []*types.User `json:"users" gorm:"many2many:user_user_group;"`
}

func userGroupList(c *gin.Context) {
	q := c.DefaultQuery("q", "")

	groups := []*userGroup{}
	db := mydb.Table("user_group")
	if len(q) > 0 {
		q = fmt.Sprintf("%%%s%%", q)
		db = db.Where("name like ?", q)
	}
	db.Find(&groups)
	for _, group := range groups {
		mydb.Model(&group).Association("Users").Find(&group.Users)
	}

	c.JSON(http.StatusOK, gin.H{"code": http.StatusOK, "groups": groups})
}

func userGroupCreate(c *gin.Context) {
	group := userGroup{}
	response := gin.H{}
	defer c.JSON(http.StatusOK, response)
	if err := c.BindJSON(&group); err != nil {
		response["code"] = http.StatusBadRequest
		response["message"] = err.Error()
		return
	}
	var cnt int
	if mydb.Table("user_group").Where("name=?", group.Name).Count(&cnt); cnt != 0 {
		response["code"] = http.StatusBadRequest
		response["message"] = "user group already exists"
		return
	}
	if err := mydb.Set("gorm:save_associations", false).Save(&group).Error; err != nil {
		response["code"] = http.StatusInternalServerError
		response["message"] = err.Error()
		return
	}
	mydb.Model(&group).Association("Users").Replace(group.Users).Find(&group.Users)

	response["code"] = http.StatusOK
	response["message"] = "group create successful"
	response["group"] = group
}

func userGroupUpdate(c *gin.Context) {
	groups := []*userGroup{}
	response := gin.H{}
	defer c.JSON(http.StatusOK, response)
	if err := c.BindJSON(&groups); err != nil {
		response["code"] = http.StatusBadRequest
		response["message"] = err.Error()
		return
	}
	var message string
	for _, group := range groups {

		if mydb.Table("user_group").Where("id=?", group.ID).First(&group).RecordNotFound() {
			response["code"] = http.StatusBadRequest
			message += "user group not found\n"
			continue
		}
		if err := mydb.Set("gorm:save_associations", false).Save(&group).Association("Users").Replace(group.Users).Error; err != nil {
			response["code"] = http.StatusInternalServerError
			message += err.Error()
			continue
		}
	}
	response["message"] = message
	response["code"] = http.StatusOK
	response["groups"] = groups
	response["message"] = "user group update successful"
}

func userGroupDelete(c *gin.Context) {
	id, _ := strconv.Atoi(c.Param("id"))
	response := gin.H{}
	defer c.JSON(http.StatusOK, response)
	group := userGroup{}
	if mydb.Table("user_group").Where("id = ?", id).Find(&group).RecordNotFound() {
		response["code"] = http.StatusNotFound
		response["message"] = "The usergroup not found"
		return
	}
	if err := mydb.Delete(&group).Error; err != nil {
		response["code"] = http.StatusInternalServerError
		response["message"] = err.Error()
		return

	}
	response["code"] = http.StatusOK
	response["message"] = fmt.Sprintf("%s delete successful", group.Name)
}
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"sync"
	"time"

	"github.com/dustin/go-humanize"
)

var (
	verbose     = flag.Bool("v", false, "verbosity")
	concurrency = flag.Int("j", 2,
		"number of concurrent dumps")
	dbName = flag.String("db", "", "which db to dump (default: all)")
)

func init() {
	log.SetFlags(log.Lmicroseconds)
}

func maybeFatal(err error, fmt string, args ...interface{}) {
	if err != nil {
		log.Fatalf(fmt, args...)
	}
}

func listDatabases(u url.URL) []string {
	u.Path = "/_all_dbs"
	res, err := http.Get(u.String())
	maybeFatal(err, "Error listing databases: %v", err)
	defer res.Body.Close()

	rv := []string{}
	d := json.NewDecoder(res.Body)
	err = d.Decode(&rv)
	maybeFatal(err, "Error decoding database list: %v", err)
	return rv
}

func vlog(s string, a ...interface{}) {
	if *verbose {
		log.Printf(s, a...)
	}
}

func compress(w io.Writer) io.WriteCloser {
	z, err := gzip.NewWriterLevel(w, gzip.BestCompression)
	maybeFatal(err, "NewWriterLevel: %v", err)
	return z
}

func dumpOne(dbname, u string) (int64, error) {
	res, err := http.Get(u)
	if err != nil {
		return 0, err
	}
	defer res.Body.Close()

	if res.StatusCode != 200 {
		return 0, fmt.Errorf("HTTP Error: %v", res.Status)
	}

	outf, err := os.Create(dbname + ".json.gz")
	if err != nil {
		return 0, err
	}
	defer outf.Close()

	z := compress(outf)
	defer z.Close()
	return io.Copy(z, res.Body)
}

func dump(wg *sync.WaitGroup, u url.URL, ch <-chan string) {
	defer wg.Done()

	for db := range ch {
		start := time.Now()
		vlog("Dumping %v", db)
		u.Path = "/" + db + "/_dump"
		n, err := dumpOne(db, u.String())
		maybeFatal(err, "Error dumping %v: %v", u.String(), err)

		vlog("Dumped %v of %v in %v",
			humanize.Bytes(uint64(n)), db, time.Since(start))
	}
}

func main() {
	flag.Parse()

	if flag.NArg() == 0 {
		log.Fatalf("Seriesly URL required")
	}

	u, err := url.Parse(flag.Arg(0))
	maybeFatal(err, "Parsing %v: %v", flag.Arg(0), err)

	wg := &sync.WaitGroup{}
	ch := make(chan string)

	for i := 0; i < *concurrency; i++ {
		wg.Add(1)
		go dump(wg, *u, ch)
	}

	if *dbName == "" {
		for _, db := range listDatabases(*u) {
			ch <- db
		}
	} else {
		ch <- *dbName
	}
	close(ch)

	wg.Wait()
}
<commit_msg>Added noop mode for dump<commit_after>package main

import (
	"compress/gzip"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"sync"
	"time"

	"github.com/dustin/go-humanize"
)

var (
	verbose     = flag.Bool("v", false, "verbosity")
	concurrency = flag.Int("j", 2,
		"number of concurrent dumps")
	dbName = flag.String("db", "", "which db to dump (default: all)")
	noop   = flag.Bool("n", false, "if true, don't actually write dumps")
)

func init() {
	log.SetFlags(log.Lmicroseconds)
}

func maybeFatal(err error, fmt string, args ...interface{}) {
	if err != nil {
		log.Fatalf(fmt, args...)
	}
}

func listDatabases(u url.URL) []string {
	u.Path = "/_all_dbs"
	res, err := http.Get(u.String())
	maybeFatal(err, "Error listing databases: %v", err)
	defer res.Body.Close()

	rv := []string{}
	d := json.NewDecoder(res.Body)
	err = d.Decode(&rv)
	maybeFatal(err, "Error decoding database list: %v", err)
	return rv
}

func vlog(s string, a ...interface{}) {
	if *verbose {
		log.Printf(s, a...)
	}
}

func compress(w io.Writer) io.WriteCloser {
	z, err := gzip.NewWriterLevel(w, gzip.BestCompression)
	maybeFatal(err, "NewWriterLevel: %v", err)
	return z
}

func dumpOne(dbname, u string) (int64, error) {
	if *noop {
		return 0, nil
	}
	res, err := http.Get(u)
	if err != nil {
		return 0, err
	}
	defer res.Body.Close()

	if res.StatusCode != 200 {
		return 0, fmt.Errorf("HTTP Error: %v", res.Status)
	}

	outf, err := os.Create(dbname + ".json.gz")
	if err != nil {
		return 0, err
	}
	defer outf.Close()

	z := compress(outf)
	defer z.Close()
	return io.Copy(z, res.Body)
}

func dump(wg *sync.WaitGroup, u url.URL, ch <-chan string) {
	defer wg.Done()

	for db := range ch {
		start := time.Now()
		vlog("Dumping %v", db)
		u.Path = "/" + db + "/_dump"
		n, err := dumpOne(db, u.String())
		maybeFatal(err, "Error dumping %v: %v", u.String(), err)

		if !*noop {
			vlog("Dumped %v of %v in %v",
				humanize.Bytes(uint64(n)), db, time.Since(start))
		}
	}
}

func main() {
	flag.Parse()

	if flag.NArg() == 0 {
		log.Fatalf("Seriesly URL required")
	}

	u, err := url.Parse(flag.Arg(0))
	maybeFatal(err, "Parsing %v: %v", flag.Arg(0), err)

	wg := &sync.WaitGroup{}
	ch := make(chan string)

	for i := 0; i < *concurrency; i++ {
		wg.Add(1)
		go dump(wg, *u, ch)
	}

	if *dbName == "" {
		for _, db := range listDatabases(*u) {
			ch <- db
		}
	} else {
		ch <- *dbName
	}
	close(ch)

	wg.Wait()
}
<|endoftext|>


<commit_before>package dnssd

import (
	"context"
	"fmt"
	"github.com/brutella/dnssd/log"
	"github.com/miekg/dns"
	"math/rand"
	"net"
	"strings"
	"time"
)

// ProbeService probes for the hostname and service instance name of srv.
// If err == nil, the returned service is verified to be unique on the local network.
func ProbeService(ctx context.Context, srv Service) (Service, error) {
	conn, err := newMDNSConn()

	if err != nil {
		return srv, err
	}

	defer conn.close()

	// After one minute of probing, if the Multicast DNS responder has been
	// unable to find any unused name, it should log an error (RFC6762 9)
	probeCtx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	// When ready to send its Multicast DNS probe packet(s) the host should
	// first wait for a short random delay time, uniformly distributed in
	// the range 0-250 ms. (RFC6762 8.1)
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	delay := time.Duration(r.Intn(250)) * time.Millisecond
	log.Debug.Println("Probing delay", delay)
	time.Sleep(delay)

	return probeService(probeCtx, conn, srv, 1*time.Millisecond, false)
}

func ReprobeService(ctx context.Context, srv Service) (Service, error) {
	conn, err := newMDNSConn()

	if err != nil {
		return srv, err
	}

	defer conn.close()
	return probeService(ctx, conn, srv, 1*time.Millisecond, true)
}

func probeService(ctx context.Context, conn MDNSConn, srv Service, delay time.Duration, probeOnce bool) (s Service, e error) {
	candidate := srv.Copy()
	prevConflict := probeConflict{}

	// Keep track of the number of conflicts
	numHostConflicts := 0
	numNameConflicts := 0

	for i := 1; i <= 100; i++ {
		conflict, err := probe(ctx, conn, *candidate)
		if err != nil {
			e = err
			return
		}

		if conflict.hasNone() {
			s = *candidate
			return
		}

		candidate = candidate.Copy()

		if conflict.hostname && (prevConflict.hostname || probeOnce) {
			numHostConflicts++
			candidate.Host = fmt.Sprintf("%s-%d", srv.Host, numHostConflicts+1)
			conflict.hostname = false
		}

		if conflict.serviceName && (prevConflict.serviceName || probeOnce) {
			numNameConflicts++
			candidate.Name = fmt.Sprintf("%s-%d", srv.Name, numNameConflicts+1)
			conflict.serviceName = false
		}

		prevConflict = conflict

		if conflict.hasAny() {
			// If the host finds that its own data is lexicographically earlier,
			// then it defers to the winning host by waiting one second,
			// and then begins probing for this record again. (RFC6762 8.2)
			log.Debug.Println("Increase wait time after receiving conflicting data")
			delay = 1 * time.Second
		} else {
			delay = 250 * time.Millisecond
		}

		log.Debug.Println("Probing wait", delay)
		time.Sleep(delay)
	}

	return
}

func probe(ctx context.Context, conn MDNSConn, service Service) (conflict probeConflict, err error) {
	for ifname, ips := range service.IfaceIPs {
		iface, err := net.InterfaceByName(ifname)
		if err != nil {
			log.Debug.Printf("error getting interface with name %s: %v\n", ifname, err)
			continue
		}
		log.Debug.Printf("Probing with %v at %s\n", ips, iface.Name)

		conflict, err := probeAtInterface(ctx, conn, service, iface)
		if !conflict.hasNone() {
			return conflict, err
		}
	}

	return probeConflict{}, nil
}

func probeAtInterface(ctx context.Context, conn MDNSConn, service Service, iface *net.Interface) (conflict probeConflict, err error) {

	msg := new(dns.Msg)

	instanceQ := dns.Question{
		Name:   service.ServiceInstanceName(),
		Qtype:  dns.TypeANY,
		Qclass: dns.ClassINET,
	}

	hostQ := dns.Question{
		Name:   service.Hostname(),
		Qtype:  dns.TypeANY,
		Qclass: dns.ClassINET,
	}

	// Responses to probe should be unicast
	setQuestionUnicast(&instanceQ)
	setQuestionUnicast(&hostQ)

	msg.Question = []dns.Question{instanceQ, hostQ}

	srv := SRV(service)
	as := A(service, iface)
	aaaas := AAAA(service, iface)

	var authority = []dns.RR{srv}
	for _, a := range as {
		authority = append(authority, a)
	}
	for _, aaaa := range aaaas {
		authority = append(authority, aaaa)
	}
	msg.Ns = authority

	readCtx, readCancel := context.WithCancel(ctx)
	defer readCancel()

	// Multicast DNS responses received *before* the first probe packet is sent
	// MUST be silently ignored. (RFC6762 8.1)
	conn.Drain(readCtx)
	ch := conn.Read(readCtx)

	queryTime := time.After(1 * time.Millisecond)
	queriesCount := 1

	for {
		select {
		case req := <-ch:
			answers := allRecords(req.msg)
			for _, answer := range answers {
				switch rr := answer.(type) {
				case *dns.A:
					for _, a := range as {
						if isDenyingA(rr, a) {
							log.Debug.Printf("%v:%d@%s denies A\n", req.from.IP, req.from.Port, req.iface.Name)
							conflict.hostname = true
							break
						}
					}

				case *dns.AAAA:
					for _, aaaa := range aaaas {
						if isDenyingAAAA(rr, aaaa) {
							log.Debug.Printf("%v:%d@%s denies AAAA\n", req.from.IP, req.from.Port, req.iface.Name)
							conflict.hostname = true
							break
						}
					}

				case *dns.SRV:
					if isDenyingSRV(rr, srv) {
						conflict.serviceName = true
					}

				default:
					break
				}
			}

		case <-ctx.Done():
			err = ctx.Err()
			return

		case <-queryTime:
			// Stop on conflict
			if !conflict.hasNone() {
				return
			}

			// Stop after 3 probe queries
			if queriesCount > 3 {
				return
			}

			queriesCount++
			log.Debug.Println("Sending probe", msg)
			q := &Query{msg: msg, iface: iface}
			conn.SendQuery(q)

			delay := 250 * time.Millisecond
			log.Debug.Println("Waiting for conflicting data", delay)
			queryTime = time.After(delay)
		}
	}

	return
}

type probeConflict struct {
	hostname    bool
	serviceName bool
}

func (pr probeConflict) hasNone() bool {
	return !pr.hostname && !pr.serviceName
}

func (pr probeConflict) hasAny() bool {
	return pr.hostname || pr.serviceName
}

func isLexicographicalLaterA(this *dns.A, that *dns.A) bool {
	if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
		log.Debug.Println("Conflicting hosts")
		if !isValidRR(this) {
			log.Debug.Println("Invalid record produces conflict")
			return true
		}

		switch compareIP(this.A.To4(), that.A.To4()) {
		case -1:
			log.Debug.Println("Lexicographical earlier")
			break
		case 1:
			log.Debug.Println("Lexicographical later")
			return true
		default:
			log.Debug.Println("Tiebreak")
			break
		}
	}

	return false
}

func isDenyingA(this *dns.A, that *dns.A) bool {
	if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
		log.Debug.Println("Conflicting hosts")

		if !isValidRR(this) {
			log.Debug.Println("Invalid record produces conflict")
			return true
		}

		switch compareIP(this.A.To4(), that.A.To4()) {
		case -1:
			log.Debug.Println("Lexicographical earlier")
			break
		case 1:
			log.Debug.Println("Lexicographical later")
			return true
		default:
			log.Debug.Println("Tiebreak")
			break
		}
	}

	return false
}

func isLexicographicalLaterAAAA(this *dns.AAAA, that *dns.AAAA) bool {
	if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
		log.Debug.Println("Conflicting hosts")
		if !isValidRR(this) {
			log.Debug.Println("Invalid record produces conflict")
			return true
		}

		switch compareIP(this.AAAA.To16(), that.AAAA.To16()) {
		case -1:
			log.Debug.Println("%s is lexicographical earlier than %s", this.AAAA.To16(), that.AAAA.To16())
			break
		case 1:
			log.Debug.Println("%s is lexicographical later than %s", this.AAAA.To16(), that.AAAA.To16())
			return true
		default:
			log.Debug.Println("Tiebreak")
			break
		}
	}

	return false
}

// isDenyingAAAA returns true if this denies that.
func isDenyingAAAA(this *dns.AAAA, that *dns.AAAA) bool {
	if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
		log.Debug.Println("Conflicting hosts")
		if !isValidRR(this) {
			log.Debug.Println("Invalid record produces conflict")
			return true
		}

		switch compareIP(this.AAAA.To16(), that.AAAA.To16()) {
		case -1:
			log.Debug.Println("%s is lexicographical earlier than %s", this.AAAA.To16(), that.AAAA.To16())
			break
		case 1:
			log.Debug.Println("%s is lexicographical later than %s", this.AAAA.To16(), that.AAAA.To16())
			return true
		default:
			log.Debug.Println("Tiebreak")
			break
		}
	}

	return false
}
func isLexicographicalLaterSRV(this *dns.SRV, that *dns.SRV) bool {
	if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
		log.Debug.Println("Conflicting SRV")
		if !isValidRR(this) {
			log.Debug.Println("Invalid record produces conflict")
			return true
		}

		switch compareSRV(this, that) {
		case -1:
			log.Debug.Println("Lexicographical earlier")
			break
		case 1:
			log.Debug.Println("Lexicographical later")
			return true
		default:
			log.Debug.Println("Tiebreak")
			break
		}
	}

	return false
}

// isDenyingSRV returns true if this denies that.
func isDenyingSRV(this *dns.SRV, that *dns.SRV) bool {
	if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
		log.Debug.Println("Conflicting SRV")
		if !isValidRR(this) {
			log.Debug.Println("Invalid record produces conflict")
			return true
		}

		switch compareSRV(this, that) {
		case -1:
			log.Debug.Println("Lexicographical earlier")
			break
		case 1:
			log.Debug.Println("Lexicographical later")
			return true
		default:
			log.Debug.Println("Tiebreak")
			break
		}
	}

	return false
}

func isValidRR(rr dns.RR) bool {
	switch r := rr.(type) {
	case *dns.A:
		return !net.IPv4zero.Equal(r.A)
	case *dns.AAAA:
		return !net.IPv6zero.Equal(r.AAAA)
	case *dns.SRV:
		return len(r.Target) > 0 && r.Port != 0
	default:
		break
	}

	return true
}

func compareIP(this net.IP, that net.IP) int {
	count := len(this)
	if count > len(that) {
		count = len(that)
	}

	for i := 0; i < count; i++ {
		if this[i] < that[i] {
			return -1
		} else if this[i] > that[i] {
			return 1
		}
	}

	if len(this) < len(that) {
		return -1
	} else if len(this) > len(that) {
		return 1
	}
	return 0
}

func compareSRV(this *dns.SRV, that *dns.SRV) int {
	if this.Priority < that.Priority {
		return -1
	} else if this.Priority > that.Priority {
		return 1
	}

	if this.Weight < that.Weight {
		return -1
	} else if this.Weight > that.Weight {
		return 1
	}

	if this.Port < that.Port {
		return -1
	} else if this.Port > that.Port {
		return 1
	}

	return strings.Compare(this.Target, that.Target)
}
<commit_msg>Refactoring<commit_after>package dnssd

import (
	"context"
	"fmt"
	"github.com/brutella/dnssd/log"
	"github.com/miekg/dns"
	"math/rand"
	"net"
	"strings"
	"time"
)

// ProbeService probes for the hostname and service instance name of srv.
// If err == nil, the returned service is verified to be unique on the local network.
func ProbeService(ctx context.Context, srv Service) (Service, error) {
	conn, err := newMDNSConn()

	if err != nil {
		return srv, err
	}

	defer conn.close()

	// After one minute of probing, if the Multicast DNS responder has been
	// unable to find any unused name, it should log an error (RFC6762 9)
	probeCtx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	// When ready to send its Multicast DNS probe packet(s) the host should
	// first wait for a short random delay time, uniformly distributed in
	// the range 0-250 ms. (RFC6762 8.1)
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	delay := time.Duration(r.Intn(250)) * time.Millisecond
	log.Debug.Println("Probing delay", delay)
	time.Sleep(delay)

	return probeService(probeCtx, conn, srv, 1*time.Millisecond, false)
}

func ReprobeService(ctx context.Context, srv Service) (Service, error) {
	conn, err := newMDNSConn()

	if err != nil {
		return srv, err
	}

	defer conn.close()
	return probeService(ctx, conn, srv, 1*time.Millisecond, true)
}

func probeService(ctx context.Context, conn MDNSConn, srv Service, delay time.Duration, probeOnce bool) (s Service, e error) {
	candidate := srv.Copy()
	prevConflict := probeConflict{}

	// Keep track of the number of conflicts
	numHostConflicts := 0
	numNameConflicts := 0

	for i := 1; i <= 100; i++ {
		conflict, err := probe(ctx, conn, *candidate)
		if err != nil {
			e = err
			return
		}

		if conflict.hasNone() {
			s = *candidate
			return
		}

		candidate = candidate.Copy()

		if conflict.hostname && (prevConflict.hostname || probeOnce) {
			numHostConflicts++
			candidate.Host = fmt.Sprintf("%s-%d", srv.Host, numHostConflicts+1)
			conflict.hostname = false
		}

		if conflict.serviceName && (prevConflict.serviceName || probeOnce) {
			numNameConflicts++
			candidate.Name = fmt.Sprintf("%s-%d", srv.Name, numNameConflicts+1)
			conflict.serviceName = false
		}

		prevConflict = conflict

		if conflict.hasAny() {
			// If the host finds that its own data is lexicographically earlier,
			// then it defers to the winning host by waiting one second,
			// and then begins probing for this record again. (RFC6762 8.2)
			log.Debug.Println("Increase wait time after receiving conflicting data")
			delay = 1 * time.Second
		} else {
			delay = 250 * time.Millisecond
		}

		log.Debug.Println("Probing wait", delay)
		time.Sleep(delay)
	}

	return
}

func probe(ctx context.Context, conn MDNSConn, service Service) (conflict probeConflict, err error) {
	for ifname, ips := range service.IfaceIPs {
		iface, err := net.InterfaceByName(ifname)
		if err != nil {
			log.Debug.Printf("error getting interface with name %s: %v\n", ifname, err)
			continue
		}
		log.Debug.Printf("Probing with %v at %s\n", ips, iface.Name)

		conflict, err := probeAtInterface(ctx, conn, service, iface)
		if conflict.hasAny() {
			return conflict, err
		}
	}

	return probeConflict{}, nil
}

func probeAtInterface(ctx context.Context, conn MDNSConn, service Service, iface *net.Interface) (conflict probeConflict, err error) {

	msg := new(dns.Msg)

	instanceQ := dns.Question{
		Name:   service.ServiceInstanceName(),
		Qtype:  dns.TypeANY,
		Qclass: dns.ClassINET,
	}

	hostQ := dns.Question{
		Name:   service.Hostname(),
		Qtype:  dns.TypeANY,
		Qclass: dns.ClassINET,
	}

	// Responses to probe should be unicast
	setQuestionUnicast(&instanceQ)
	setQuestionUnicast(&hostQ)

	msg.Question = []dns.Question{instanceQ, hostQ}

	srv := SRV(service)
	as := A(service, iface)
	aaaas := AAAA(service, iface)

	var authority = []dns.RR{srv}
	for _, a := range as {
		authority = append(authority, a)
	}
	for _, aaaa := range aaaas {
		authority = append(authority, aaaa)
	}
	msg.Ns = authority

	readCtx, readCancel := context.WithCancel(ctx)
	defer readCancel()

	// Multicast DNS responses received *before* the first probe packet is sent
	// MUST be silently ignored. (RFC6762 8.1)
	conn.Drain(readCtx)
	ch := conn.Read(readCtx)

	queryTime := time.After(1 * time.Millisecond)
	queriesCount := 1

	for {
		select {
		case req := <-ch:
			answers := allRecords(req.msg)
			for _, answer := range answers {
				switch rr := answer.(type) {
				case *dns.A:
					for _, a := range as {
						if isDenyingA(rr, a) {
							log.Debug.Printf("%v:%d@%s denies A\n", req.from.IP, req.from.Port, req.iface.Name)
							conflict.hostname = true
							break
						}
					}

				case *dns.AAAA:
					for _, aaaa := range aaaas {
						if isDenyingAAAA(rr, aaaa) {
							log.Debug.Printf("%v:%d@%s denies AAAA\n", req.from.IP, req.from.Port, req.iface.Name)
							conflict.hostname = true
							break
						}
					}

				case *dns.SRV:
					if isDenyingSRV(rr, srv) {
						conflict.serviceName = true
					}

				default:
					break
				}
			}

		case <-ctx.Done():
			err = ctx.Err()
			return

		case <-queryTime:
			// Stop on conflict
			if conflict.hasAny() {
				return
			}

			// Stop after 3 probe queries
			if queriesCount > 3 {
				return
			}

			queriesCount++
			log.Debug.Println("Sending probe", msg)
			q := &Query{msg: msg, iface: iface}
			conn.SendQuery(q)

			delay := 250 * time.Millisecond
			log.Debug.Println("Waiting for conflicting data", delay)
			queryTime = time.After(delay)
		}
	}

	return
}

type probeConflict struct {
	hostname    bool
	serviceName bool
}

func (pr probeConflict) hasNone() bool {
	return !pr.hostname && !pr.serviceName
}

func (pr probeConflict) hasAny() bool {
	return pr.hostname || pr.serviceName
}

func isDenyingA(this *dns.A, that *dns.A) bool {
	if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
		log.Debug.Println("Conflicting hosts")

		if !isValidRR(this) {
			log.Debug.Println("Invalid record produces conflict")
			return true
		}

		switch compareIP(this.A.To4(), that.A.To4()) {
		case -1:
			log.Debug.Println("Lexicographical earlier")
			break
		case 1:
			log.Debug.Println("Lexicographical later")
			return true
		default:
			log.Debug.Println("Tiebreak")
			break
		}
	}

	return false
}

// isDenyingAAAA returns true if this denies that.
func isDenyingAAAA(this *dns.AAAA, that *dns.AAAA) bool {
	if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
		log.Debug.Println("Conflicting hosts")
		if !isValidRR(this) {
			log.Debug.Println("Invalid record produces conflict")
			return true
		}

		switch compareIP(this.AAAA.To16(), that.AAAA.To16()) {
		case -1:
			log.Debug.Println("%s is lexicographical earlier than %s", this.AAAA.To16(), that.AAAA.To16())
			break
		case 1:
			log.Debug.Println("%s is lexicographical later than %s", this.AAAA.To16(), that.AAAA.To16())
			return true
		default:
			log.Debug.Println("Tiebreak")
			break
		}
	}

	return false
}

// isDenyingSRV returns true if this denies that.
func isDenyingSRV(this *dns.SRV, that *dns.SRV) bool {
	if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
		log.Debug.Println("Conflicting SRV")
		if !isValidRR(this) {
			log.Debug.Println("Invalid record produces conflict")
			return true
		}

		switch compareSRV(this, that) {
		case -1:
			log.Debug.Println("Lexicographical earlier")
			break
		case 1:
			log.Debug.Println("Lexicographical later")
			return true
		default:
			log.Debug.Println("Tiebreak")
			break
		}
	}

	return false
}

func isValidRR(rr dns.RR) bool {
	switch r := rr.(type) {
	case *dns.A:
		return !net.IPv4zero.Equal(r.A)
	case *dns.AAAA:
		return !net.IPv6zero.Equal(r.AAAA)
	case *dns.SRV:
		return len(r.Target) > 0 && r.Port != 0
	default:
		break
	}

	return true
}

func compareIP(this net.IP, that net.IP) int {
	count := len(this)
	if count > len(that) {
		count = len(that)
	}

	for i := 0; i < count; i++ {
		if this[i] < that[i] {
			return -1
		} else if this[i] > that[i] {
			return 1
		}
	}

	if len(this) < len(that) {
		return -1
	} else if len(this) > len(that) {
		return 1
	}
	return 0
}

func compareSRV(this *dns.SRV, that *dns.SRV) int {
	if this.Priority < that.Priority {
		return -1
	} else if this.Priority > that.Priority {
		return 1
	}

	if this.Weight < that.Weight {
		return -1
	} else if this.Weight > that.Weight {
		return 1
	}

	if this.Port < that.Port {
		return -1
	} else if this.Port > that.Port {
		return 1
	}

	return strings.Compare(this.Target, that.Target)
}
<|endoftext|>


<commit_before>package sql

import (
	"database/sql"
	"errors"
	"fmt"
	"reflect"
	"strconv"
)

func NewPool(sourceStrings []string) (pool *Pool, err error) {
	queue := make(chan *sql.DB, len(sourceStrings))
	var conn *sql.DB
	for _, sourceStr := range sourceStrings {
		conn, err = sql.Open("mysql", sourceStr)
		if err != nil {
			return
		}
		err = conn.Ping()
		if err != nil {
			return
		}

		queue <- conn
	}
	pool = &Pool{queue}
	return
}

type Pool struct {
	queue chan *sql.DB
}

func (p *Pool) Autocommit(acFun func(ac *Pool) error) (err error) {
	conn := <-p.queue
	defer func() { p.queue <- conn }()

	acPool := &Pool{queue: make(chan *sql.DB, 1)}
	acPool.queue <- conn

	return acFun(acPool)
}

func (p *Pool) Transact(txFun func(tx *Pool) error) (err error) {
	conn := <-p.queue
	defer func() { p.queue <- conn }()

	txPool := &Pool{queue: make(chan *sql.DB, 1)}
	txPool.queue <- conn

	_, err = conn.Exec("START TRANSACTION")
	if err != nil {
		return
	}

	err = txFun(txPool)
	if err != nil {
		return rollback(conn, err)
	}

	_, err = conn.Exec("COMMIT")
	if err != nil {
		return
	}

	return
}

func rollback(conn *sql.DB, err error) error {
	_, rbErr := conn.Exec("ROLLBACK")
	if rbErr != nil {
		err = errors.New("Rollback error: " + rbErr.Error() + " Query error: " + err.Error())
	}
	return err
}

// Query with fixed args
func (p *Pool) Query(query string, args ...interface{}) (rows *sql.Rows, err error) {
	conn := <-p.queue
	defer func() { p.queue <- conn }()

	FixArgs(args)
	rows, err = conn.Query(query, args...)
	if err != nil {
		err = errors.New("sql.Query Error: " + err.Error() + ". Query: " + query + " Args: " + fmt.Sprint(args))
	}
	return
}

// Execute with fixed args
func (p *Pool) Exec(query string, args ...interface{}) (res sql.Result, err error) {
	conn := <-p.queue
	defer func() { p.queue <- conn }()

	FixArgs(args)
	res, err = conn.Exec(query, args...)
	if err != nil {
		err = errors.New("sql.Exec Error: " + err.Error() + ". Query: " + query + " Args: " + fmt.Sprint(args))
	}
	return
}

/*
Fix args by converting them to values of their underlying kind.
This avoids problems in database/sql with e.g custom string types.
Without FixArgs, the following code:

	type Foo string
	...
	pool.Query("SELECT * WHERE Foo=?", Foo("bar"))

would give you the error:

	sql: converting Exec argument #1's type: unsupported type Foo, a string
*/
func FixArgs(args []interface{}) {
	for i, arg := range args {
		vArg := reflect.ValueOf(arg)
		if vArg.Kind() == reflect.String {
			args[i] = vArg.String()
		}
	}
}

func (p *Pool) SelectInt(query string, args ...interface{}) (num int, err error) {
	err = p.queryOne(query, args, &num)
	return
}

func (p *Pool) SelectString(query string, args ...interface{}) (str string, err error) {
	err = p.queryOne(query, args, &str)
	return
}

func (p *Pool) SelectUInt(query string, args ...interface{}) (num uint, err error) {
	err = p.queryOne(query, args, &num)
	return
}

func (p *Pool) queryOne(query string, args []interface{}, out interface{}) (err error) {
	rows, err := p.Query(query, args...)
	if err != nil {
		return
	}

	if rows.Next() {
		err = rows.Scan(out)
		if err != nil {
			return
		}
	} else {
		err = errors.New("Query returned no rows")
		return
	}

	if rows.Next() {
		err = errors.New("Query returned too many rows")
		return
	}

	return
}

// type TransactFunc func() error

// func (p *Pool) Transact(f TransactFunc) (err error) {
// 	conn := <-p.queue
// 	defer func() { p.queue <- conn }()

// 	_, err = conn.Exec("START TRANSACTION")
// 	if err != nil {
// 		return
// 	}

// 	err = f(conn)
// 	if err != nil {
// 		_, rollbackError := conn.Exec("ROLLBACK")
// 		if rollbackError != nil {
// 			panic("Could not rollback transaction. TransactionFunc error: " + err.Error() + ". Rollback error:" + rollbackError.Error())
// 		}

// 	} else {
// 		_, err = conn.Exec("COMMIT")
// 	}

// 	return
// }

func (p *Pool) UpdateOne(query string, args ...interface{}) error {
	res, err := p.Exec(query, args...)
	if err != nil {
		return err
	}

	affected, err := res.RowsAffected()
	if err != nil {
		return err
	}
	if affected != 1 {
		return errors.New(fmt.Sprintf("UpdateOne affected %d rows. Query: %q Args: %q", affected, query, args))
	}

	return nil
}

func (p *Pool) InsertIgnoreId(query string, args ...interface{}) (err error) {
	_, err = p.Insert(query, args...)
	return
}

func (p *Pool) Insert(query string, args ...interface{}) (id int64, err error) {
	res, err := p.Exec(query, args...)
	if err != nil {
		return
	}
	id, err = res.LastInsertId()
	return
}

func (p *Pool) Select(output interface{}, sql string, args ...interface{}) error {
	// Check types
	var outputPtr = reflect.ValueOf(output)
	if outputPtr.Kind() != reflect.Ptr {
		return errors.New("fun/sql.Select: expects a pointer to a slice of items")
	}
	var outputReflection = reflect.Indirect(outputPtr)
	if outputReflection.Kind() != reflect.Slice {
		return errors.New("fun/sql.Select: expects items to be a slice")
	}
	if outputReflection.Len() != 0 {
		return errors.New("fun/sql.Select: expects items to be empty")
	}
	outputItemType := outputReflection.Type().Elem().Elem()
	if outputItemType.Kind() != reflect.Struct {
		return errors.New("fun/sql.Select: expects items to be a slice of structs")
	}

	// Query DB
	var rows, err = p.Query(sql, args...)
	if err != nil {
		return err
	}

	// Reflect onto structs
	columns, err := rows.Columns()
	if err != nil {
		return err
	}

	structType := outputReflection.Type().Elem()
	for rows.Next() {
		structPtrVal := reflect.New(structType.Elem())
		outputItemStructVal := structPtrVal.Elem()
		err = structFromRow(outputItemStructVal, columns, rows)
		if err != nil {
			return err
		}

		outputReflection.Set(reflect.Append(outputReflection, structPtrVal))
	}
	return nil
}

const selectOneTypeError = "fun/sql.SelectOne: expects a **struct, e.g var person *Person; c.SelectOne(&person, sql)"

func (p *Pool) SelectOne(output interface{}, query string, args ...interface{}) error {
	return p.selectOne(output, query, true, args...)
}
func (p *Pool) SelectMaybe(output interface{}, query string, args ...interface{}) error {
	return p.selectOne(output, query, false, args...)
}
func (p *Pool) selectOne(output interface{}, query string, required bool, args ...interface{}) error {
	// Check types
	var outputReflectionPtr = reflect.ValueOf(output)
	if !outputReflectionPtr.IsValid() {
		panic(selectOneTypeError)
	}
	if outputReflectionPtr.Kind() != reflect.Ptr {
		panic(selectOneTypeError)
	}
	var outputReflection = outputReflectionPtr.Elem()
	if outputReflection.Kind() != reflect.Ptr {
		panic(selectOneTypeError)
	}

	// Query DB
	var rows, err = p.Query(query, args...)
	if err != nil {
		return err
	}

	// Reflect onto struct
	columns, err := rows.Columns()
	if err != nil {
		return err
	}
	if !rows.Next() {
		if required {
			return errors.New("SelectOne got 0 rows. Query: " + query)
		} else {
			return nil
		}
	}

	var vStruct reflect.Value
	if outputReflection.IsNil() {
		structPtrVal := reflect.New(outputReflection.Type().Elem())
		outputReflection.Set(structPtrVal)
		vStruct = structPtrVal.Elem()
	} else {
		vStruct = outputReflection.Elem()
	}

	err = structFromRow(vStruct, columns, rows)
	if err != nil {
		return err
	}

	if rows.Next() {
		return errors.New("fun/sql.SelectOne: got multiple rows. Query: " + query + " Args: " + fmt.Sprint(args))
	}

	return nil
}

func structFromRow(outputItemStructVal reflect.Value, columns []string, rows *sql.Rows) (err error) {
	vals := make([]interface{}, len(columns))
	for i, _ := range columns {
		vals[i] = &sql.RawBytes{}
		// vals[i] = &[]byte{}
	}
	err = rows.Scan(vals...)
	if err != nil {
		return
	}

	var uintVal uint64
	var intVal int64
	for i, column := range columns {
		bytes := []byte(*vals[i].(*sql.RawBytes))
		// bytes := []byte(*vals[i].(*[]byte))
		if bytes == nil {
			continue // Leave struct field empty
		}
		var outputItemField = outputItemStructVal.FieldByName(column)
		switch outputItemField.Kind() {
		case reflect.String:
			outputItemField.SetString(string(bytes))
		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
			uintVal, err = strconv.ParseUint(string(bytes), 10, 64)
			if err != nil {
				return
			}
			outputItemField.SetUint(reflect.ValueOf(uintVal).Uint())
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			intVal, err = strconv.ParseInt(string(bytes), 10, 64)
			if err != nil {
				return
			}
			outputItemField.SetInt(reflect.ValueOf(intVal).Int())
		default:
			err = errors.New("fun/sql: Bad row value for column: " + column + " " + outputItemField.Kind().String())
			return
		}
	}

	return
}
<commit_msg>Add Select<Int/String/Uint>Maybe<commit_after>package sql

import (
	"database/sql"
	"errors"
	"fmt"
	"reflect"
	"strconv"
)

func NewPool(sourceStrings []string) (pool *Pool, err error) {
	queue := make(chan *sql.DB, len(sourceStrings))
	var conn *sql.DB
	for _, sourceStr := range sourceStrings {
		conn, err = sql.Open("mysql", sourceStr)
		if err != nil {
			return
		}
		err = conn.Ping()
		if err != nil {
			return
		}

		queue <- conn
	}
	pool = &Pool{queue}
	return
}

type Pool struct {
	queue chan *sql.DB
}

func (p *Pool) Autocommit(acFun func(ac *Pool) error) (err error) {
	conn := <-p.queue
	defer func() { p.queue <- conn }()

	acPool := &Pool{queue: make(chan *sql.DB, 1)}
	acPool.queue <- conn

	return acFun(acPool)
}

func (p *Pool) Transact(txFun func(tx *Pool) error) (err error) {
	conn := <-p.queue
	defer func() { p.queue <- conn }()

	txPool := &Pool{queue: make(chan *sql.DB, 1)}
	txPool.queue <- conn

	_, err = conn.Exec("START TRANSACTION")
	if err != nil {
		return
	}

	err = txFun(txPool)
	if err != nil {
		return rollback(conn, err)
	}

	_, err = conn.Exec("COMMIT")
	if err != nil {
		return
	}

	return
}

func rollback(conn *sql.DB, err error) error {
	_, rbErr := conn.Exec("ROLLBACK")
	if rbErr != nil {
		err = errors.New("Rollback error: " + rbErr.Error() + " Query error: " + err.Error())
	}
	return err
}

// Query with fixed args
func (p *Pool) Query(query string, args ...interface{}) (rows *sql.Rows, err error) {
	conn := <-p.queue
	defer func() { p.queue <- conn }()

	FixArgs(args)
	rows, err = conn.Query(query, args...)
	if err != nil {
		err = errors.New("sql.Query Error: " + err.Error() + ". Query: " + query + " Args: " + fmt.Sprint(args))
	}
	return
}

// Execute with fixed args
func (p *Pool) Exec(query string, args ...interface{}) (res sql.Result, err error) {
	conn := <-p.queue
	defer func() { p.queue <- conn }()

	FixArgs(args)
	res, err = conn.Exec(query, args...)
	if err != nil {
		err = errors.New("sql.Exec Error: " + err.Error() + ". Query: " + query + " Args: " + fmt.Sprint(args))
	}
	return
}

/*
Fix args by converting them to values of their underlying kind.
This avoids problems in database/sql with e.g custom string types.
Without FixArgs, the following code:

	type Foo string
	...
	pool.Query("SELECT * WHERE Foo=?", Foo("bar"))

would give you the error:

	sql: converting Exec argument #1's type: unsupported type Foo, a string
*/
func FixArgs(args []interface{}) {
	for i, arg := range args {
		vArg := reflect.ValueOf(arg)
		if vArg.Kind() == reflect.String {
			args[i] = vArg.String()
		}
	}
}

func (p *Pool) SelectInt(query string, args ...interface{}) (num int, err error) {
	found, err := p.queryOne(query, args, &num)
	if !found {
		err = errors.New("Query returned no rows")
	}
	return
}

func (p *Pool) SelectString(query string, args ...interface{}) (str string, err error) {
	found, err := p.queryOne(query, args, &str)
	if !found {
		err = errors.New("Query returned no rows")
	}
	return
}

func (p *Pool) SelectUInt(query string, args ...interface{}) (num uint, err error) {
	found, err := p.queryOne(query, args, &num)
	if !found {
		err = errors.New("Query returned no rows")
	}
	return
}

func (p *Pool) SelectIntMaybe(query string, args ...interface{}) (num int, found bool, err error) {
	found, err = p.queryOne(query, args, &num)
	return
}

func (p *Pool) SelectStringMaybe(query string, args ...interface{}) (str string, found bool, err error) {
	found, err = p.queryOne(query, args, &str)
	return
}

func (p *Pool) SelectUIntMaybe(query string, args ...interface{}) (num uint, found bool, err error) {
	found, err = p.queryOne(query, args, &num)
	return
}

func (p *Pool) queryOne(query string, args []interface{}, out interface{}) (found bool, err error) {
	rows, err := p.Query(query, args...)
	if err != nil {
		return
	}

	if rows.Next() {
		err = rows.Scan(out)
		if err != nil {
			return
		}
		if rows.Next() {
			err = errors.New("Query returned too many rows")
			return
		}
		found = true
	} else {
		found = false
	}

	return
}

// type TransactFunc func() error

// func (p *Pool) Transact(f TransactFunc) (err error) {
// 	conn := <-p.queue
// 	defer func() { p.queue <- conn }()

// 	_, err = conn.Exec("START TRANSACTION")
// 	if err != nil {
// 		return
// 	}

// 	err = f(conn)
// 	if err != nil {
// 		_, rollbackError := conn.Exec("ROLLBACK")
// 		if rollbackError != nil {
// 			panic("Could not rollback transaction. TransactionFunc error: " + err.Error() + ". Rollback error:" + rollbackError.Error())
// 		}

// 	} else {
// 		_, err = conn.Exec("COMMIT")
// 	}

// 	return
// }

func (p *Pool) UpdateOne(query string, args ...interface{}) error {
	res, err := p.Exec(query, args...)
	if err != nil {
		return err
	}

	affected, err := res.RowsAffected()
	if err != nil {
		return err
	}
	if affected != 1 {
		return errors.New(fmt.Sprintf("UpdateOne affected %d rows. Query: %q Args: %q", affected, query, args))
	}

	return nil
}

func (p *Pool) InsertIgnoreId(query string, args ...interface{}) (err error) {
	_, err = p.Insert(query, args...)
	return
}

func (p *Pool) Insert(query string, args ...interface{}) (id int64, err error) {
	res, err := p.Exec(query, args...)
	if err != nil {
		return
	}
	id, err = res.LastInsertId()
	return
}

func (p *Pool) Select(output interface{}, sql string, args ...interface{}) error {
	// Check types
	var outputPtr = reflect.ValueOf(output)
	if outputPtr.Kind() != reflect.Ptr {
		return errors.New("fun/sql.Select: expects a pointer to a slice of items")
	}
	var outputReflection = reflect.Indirect(outputPtr)
	if outputReflection.Kind() != reflect.Slice {
		return errors.New("fun/sql.Select: expects items to be a slice")
	}
	if outputReflection.Len() != 0 {
		return errors.New("fun/sql.Select: expects items to be empty")
	}
	outputItemType := outputReflection.Type().Elem().Elem()
	if outputItemType.Kind() != reflect.Struct {
		return errors.New("fun/sql.Select: expects items to be a slice of structs")
	}

	// Query DB
	var rows, err = p.Query(sql, args...)
	if err != nil {
		return err
	}

	// Reflect onto structs
	columns, err := rows.Columns()
	if err != nil {
		return err
	}

	structType := outputReflection.Type().Elem()
	for rows.Next() {
		structPtrVal := reflect.New(structType.Elem())
		outputItemStructVal := structPtrVal.Elem()
		err = structFromRow(outputItemStructVal, columns, rows)
		if err != nil {
			return err
		}

		outputReflection.Set(reflect.Append(outputReflection, structPtrVal))
	}
	return nil
}

const selectOneTypeError = "fun/sql.SelectOne: expects a **struct, e.g var person *Person; c.SelectOne(&person, sql)"

func (p *Pool) SelectOne(output interface{}, query string, args ...interface{}) error {
	return p.selectOne(output, query, true, args...)
}
func (p *Pool) SelectMaybe(output interface{}, query string, args ...interface{}) error {
	return p.selectOne(output, query, false, args...)
}
func (p *Pool) selectOne(output interface{}, query string, required bool, args ...interface{}) error {
	// Check types
	var outputReflectionPtr = reflect.ValueOf(output)
	if !outputReflectionPtr.IsValid() {
		panic(selectOneTypeError)
	}
	if outputReflectionPtr.Kind() != reflect.Ptr {
		panic(selectOneTypeError)
	}
	var outputReflection = outputReflectionPtr.Elem()
	if outputReflection.Kind() != reflect.Ptr {
		panic(selectOneTypeError)
	}

	// Query DB
	var rows, err = p.Query(query, args...)
	if err != nil {
		return err
	}

	// Reflect onto struct
	columns, err := rows.Columns()
	if err != nil {
		return err
	}
	if !rows.Next() {
		if required {
			return errors.New("SelectOne got 0 rows. Query: " + query)
		} else {
			return nil
		}
	}

	var vStruct reflect.Value
	if outputReflection.IsNil() {
		structPtrVal := reflect.New(outputReflection.Type().Elem())
		outputReflection.Set(structPtrVal)
		vStruct = structPtrVal.Elem()
	} else {
		vStruct = outputReflection.Elem()
	}

	err = structFromRow(vStruct, columns, rows)
	if err != nil {
		return err
	}

	if rows.Next() {
		return errors.New("fun/sql.SelectOne: got multiple rows. Query: " + query + " Args: " + fmt.Sprint(args))
	}

	return nil
}

func structFromRow(outputItemStructVal reflect.Value, columns []string, rows *sql.Rows) (err error) {
	vals := make([]interface{}, len(columns))
	for i, _ := range columns {
		vals[i] = &sql.RawBytes{}
		// vals[i] = &[]byte{}
	}
	err = rows.Scan(vals...)
	if err != nil {
		return
	}

	var uintVal uint64
	var intVal int64
	for i, column := range columns {
		bytes := []byte(*vals[i].(*sql.RawBytes))
		// bytes := []byte(*vals[i].(*[]byte))
		if bytes == nil {
			continue // Leave struct field empty
		}
		var outputItemField = outputItemStructVal.FieldByName(column)
		switch outputItemField.Kind() {
		case reflect.String:
			outputItemField.SetString(string(bytes))
		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
			uintVal, err = strconv.ParseUint(string(bytes), 10, 64)
			if err != nil {
				return
			}
			outputItemField.SetUint(reflect.ValueOf(uintVal).Uint())
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			intVal, err = strconv.ParseInt(string(bytes), 10, 64)
			if err != nil {
				return
			}
			outputItemField.SetInt(reflect.ValueOf(intVal).Int())
		default:
			err = errors.New("fun/sql: Bad row value for column: " + column + " " + outputItemField.Kind().String())
			return
		}
	}

	return
}
<|endoftext|>


<commit_before>package gortcdc

import (
  "math/rand"
  "time"
  "fmt"
  "strings"
  "strconv"
  "encoding/json"

  log "github.com/Sirupsen/logrus"

  "github.com/xhs/goblice"
  "github.com/xhs/godtls"
  "github.com/xhs/gosctp"
)

func init() {
  log.SetLevel(log.DebugLevel)
}

const (
  roleClient = 0
  roleServer = 1

  stateClosed = 0
  stateConnecting = 1
  stateConnected = 2
)

type Peer struct {
  ice *goblice.Agent
  ctx *godtls.DtlsContext
  dtls *godtls.DtlsTransport
  sctp *gosctp.SctpTransport
  role int
  remotePort int
  state int
}

func NewPeer() (*Peer, error) {
  ice, err := goblice.NewReliableAgent()
  if err != nil {
    return nil, err
  }
  ctx, err := godtls.NewContext("gortcdc", 365)
  if err != nil {
    ice.Destroy()
    return nil, err
  }
  dtls, err := ctx.NewTransport()
  if err != nil {
    ice.Destroy()
    ctx.Destroy()
    return nil, err
  }
  rand.Seed(time.Now().UnixNano())
  sctp, err := gosctp.NewTransport(rand.Intn(50001) + 10000)
  if err != nil {
    ice.Destroy()
    ctx.Destroy()
    dtls.Destroy()
    return nil, err
  }
  p := &Peer{
    ice: ice,
    ctx: ctx,
    dtls: dtls,
    sctp: sctp,
    role: roleClient,
    state: stateClosed,
  }
  return p, nil
}

func (p *Peer) Destroy() {
  p.ice.Destroy()
  p.dtls.Destroy()
  p.ctx.Destroy()
  p.sctp.Destroy()
}

type Signaller interface {
  Send(data []byte) error
  ReceiveFrom() <-chan []byte
}

type Signal struct {
  Type string `json:"type"`
  Value string `json:"value"`
}

func (p *Peer) Run(s Signaller) error {
  if p.state == stateConnecting {
    if err := p.ice.GatherCandidates(); err != nil {
      return err
    }
  }
  go p.ice.Run()

  signalChan := s.ReceiveFrom()
  for {
    select {
    case cand := <-p.ice.CandidateChannel:
      candidate, err := json.Marshal(&Signal{"candidate", cand})
      if err != nil {
        return err
      }
      s.Send(candidate)
      continue
    case e := <-p.ice.EventChannel:
      if e == goblice.EventNegotiationDone {
        break
      }
      continue
    case data := <-signalChan:
      log.Debug("signal received:", string(data))
      sig := &Signal{}
      if err := json.Unmarshal(data, sig); err != nil {
        continue
      }
      if sig.Type == "candidate" {
        p.ParseCandidateSdp(sig.Value)
      } else if sig.Type == "offer" {
        p.ParseOfferSdp(sig.Value)
        sdp, err := p.GenerateOfferSdp()
        if err != nil {
          return err
        }
        answer, err := json.Marshal(&Signal{"answer", sdp})
        if err != nil {
          return err
        }
        s.Send(answer)
        p.ice.GatherCandidates()
      } else if sig.Type == "answer" {
        p.ParseOfferSdp(sig.Value)
      }
      continue
    }
    break
  }
  log.Debug("ICE negotiation done")

  if p.role == roleClient {
    log.Debug("DTLS connecting")
    p.dtls.SetConnectState()
  } else {
    log.Debug("DTLS accepting")
    p.dtls.SetAcceptState()
  }

  go func () {
    var buf [1 << 16]byte
    for {
      data := <-p.ice.DataChannel
      log.Debug(len(data), " bytes of DTLS data received")
      p.dtls.Feed(data)

      n, _ := p.dtls.Read(buf[:])
      if n > 0 {
        log.Debug(n, " bytes of SCTP data received")
        p.sctp.Feed(buf[:n])
      }
    }
  }()

  go func () {
    var buf [1 << 16]byte
    tick := time.Tick(4 * time.Millisecond)
    for {
      <-tick
      n, _ := p.dtls.Spew(buf[:])
      if n > 0 {
        log.Debug(n, " bytes of DTLS data ready")
        rv, err := p.ice.Send(buf[:n])
        if err != nil {
          log.Warn(err)
        } else {
          log.Debug(rv, " bytes of DTLS data sent")
        }
      }
    }
  }()

  if err := p.dtls.Handshake(); err != nil {
    return err
  }
  log.Debug("DTLS handshake done")

  go func () {
    for {
      data := <-p.sctp.BufferChannel
      log.Debug(len(data), " bytes of SCTP data ready")
      p.dtls.Write(data)
    }
  }()

  if p.role == roleClient {
    if err := p.sctp.Connect(p.remotePort); err != nil {
      return err
    }
  } else {
    if err := p.sctp.Accept(); err != nil {
      return err
    }
  }
  log.Debug("SCTP handshake done")

  return nil
}

var numbers = []rune("0123456789")

func randSession() string {
  s := make([]rune, 16)
  rand.Seed(time.Now().UnixNano())
  for i := range s {
    s[i] = numbers[rand.Intn(10)]
  }
  return string(s)
}

func (p *Peer) GenerateOfferSdp() (string, error) {
  var offer []string
  offer = append(offer, "v=0")
  offer = append(offer, fmt.Sprintf("o=- %s 2 IN IP4 127.0.0.1", randSession()))
  offer = append(offer, "s=-")
  offer = append(offer, "t=0 0")
  offer = append(offer, "a=msid-semantic: WMS")
  offer = append(offer, "m=application 1 UDP/DTLS/SCTP webrtc-datachannel")
  offer = append(offer, "c=IN IP4 0.0.0.0")

  sdp := p.ice.GenerateSdp()
  sdps := strings.Split(sdp, "\n")
  for i := range sdps {
    if strings.HasPrefix(sdps[i], "a=ice-ufrag:") {
      offer = append(offer, sdps[i])
    } else if strings.HasPrefix(sdps[i], "a=ice-pwd:") {
      offer = append(offer, sdps[i])
    }
  }

  offer = append(offer, fmt.Sprintf("a=fingerprint:sha-256 %s", p.dtls.Fingerprint))
  if p.role == roleClient {
    offer = append(offer, "a=setup:active")
  } else {
    offer = append(offer, "a=setup:passive")
  }
  offer = append(offer, "a=mid:data")
  offer = append(offer, fmt.Sprintf("a=sctp-port:%d", p.sctp.Port))
  offer = append(offer, "")

  p.state = stateConnecting

  return strings.Join(offer, "\r\n"), nil
}

func (p *Peer) ParseOfferSdp(offer string) (int, error) {
  sdps := strings.Split(offer, "\r\n")
  for i := range sdps {
    if strings.HasPrefix(sdps[i], "a=sctp-port:") {
      sctpmap := strings.Split(sdps[i], " ")[0]
      port, err := strconv.Atoi(strings.Split(sctpmap, ":")[1])
      if err != nil {
        return 0, err
      }
      p.remotePort = port
    } else if strings.HasPrefix(sdps[i], "a=setup:active") {
      if p.role == roleClient {
        p.role = roleServer
      }
    } else if strings.HasPrefix(sdps[i], "a=setup:passive") {
      if p.role == roleServer {
        p.role = roleClient
      }
    }
  }

  rv, err := p.ice.ParseSdp(strings.Join(sdps, "\n"))
  if err != nil {
    return 0, err
  }

  p.state = stateConnecting

  return int(rv), nil
}

func (p *Peer) ParseCandidateSdp(cands string) int {
  sdps := strings.Split(cands, "\r\n")
  count := 0
  for i := range sdps {
    if sdps[i] == "" {
      continue
    }
    rv, err := p.ice.ParseCandidateSdp(sdps[i])
    if err != nil {
      continue
    }
    count = count + rv
  }
  return count
}
<commit_msg>make it async, remove time.Tick<commit_after>package gortcdc

import (
  "math/rand"
  "time"
  "fmt"
  "strings"
  "strconv"
  "encoding/json"

  log "github.com/Sirupsen/logrus"

  "github.com/xhs/goblice"
  "github.com/xhs/godtls"
  "github.com/xhs/gosctp"
)

func init() {
  log.SetLevel(log.DebugLevel)
}

const (
  roleClient = 0
  roleServer = 1

  stateClosed = 0
  stateConnecting = 1
  stateConnected = 2
)

type Peer struct {
  ice *goblice.Agent
  ctx *godtls.DtlsContext
  dtls *godtls.DtlsTransport
  sctp *gosctp.SctpTransport
  role int
  remotePort int
  state int
}

func NewPeer() (*Peer, error) {
  ice, err := goblice.NewReliableAgent()
  if err != nil {
    return nil, err
  }
  ctx, err := godtls.NewContext("gortcdc", 365)
  if err != nil {
    ice.Destroy()
    return nil, err
  }
  dtls, err := ctx.NewTransport()
  if err != nil {
    ice.Destroy()
    ctx.Destroy()
    return nil, err
  }
  rand.Seed(time.Now().UnixNano())
  sctp, err := gosctp.NewTransport(rand.Intn(50001) + 10000)
  if err != nil {
    ice.Destroy()
    ctx.Destroy()
    dtls.Destroy()
    return nil, err
  }
  p := &Peer{
    ice: ice,
    ctx: ctx,
    dtls: dtls,
    sctp: sctp,
    role: roleClient,
    state: stateClosed,
  }
  return p, nil
}

func (p *Peer) Destroy() {
  p.ice.Destroy()
  p.dtls.Destroy()
  p.ctx.Destroy()
  p.sctp.Destroy()
}

type Signaller interface {
  Send(data []byte) error
  ReceiveFrom() <-chan []byte
}

type Signal struct {
  Type string `json:"type"`
  Value string `json:"value"`
}

func (p *Peer) Run(s Signaller) error {
  if p.state == stateConnecting {
    if err := p.ice.GatherCandidates(); err != nil {
      return err
    }
  }
  go p.ice.Run()

  signalChan := s.ReceiveFrom()
  for {
    select {
    case cand := <-p.ice.CandidateChannel:
      candidate, err := json.Marshal(&Signal{"candidate", cand})
      if err != nil {
        return err
      }
      s.Send(candidate)
      continue
    case e := <-p.ice.EventChannel:
      if e == goblice.EventNegotiationDone {
        break
      }
      continue
    case data := <-signalChan:
      log.Debug("signal received:", string(data))
      sig := &Signal{}
      if err := json.Unmarshal(data, sig); err != nil {
        continue
      }
      if sig.Type == "candidate" {
        p.ParseCandidateSdp(sig.Value)
      } else if sig.Type == "offer" {
        p.ParseOfferSdp(sig.Value)
        sdp, err := p.GenerateOfferSdp()
        if err != nil {
          return err
        }
        answer, err := json.Marshal(&Signal{"answer", sdp})
        if err != nil {
          return err
        }
        s.Send(answer)
        p.ice.GatherCandidates()
      } else if sig.Type == "answer" {
        p.ParseOfferSdp(sig.Value)
      }
      continue
    }
    break
  }
  log.Debug("ICE negotiation done")

  if p.role == roleClient {
    log.Debug("DTLS connecting")
    p.dtls.SetConnectState()
  } else {
    log.Debug("DTLS accepting")
    p.dtls.SetAcceptState()
  }

  go func () {
    var buf [1 << 16]byte
    for {
      data := <-p.ice.DataChannel
      log.Debug(len(data), " bytes of DTLS data received")
      p.dtls.Feed(data)

      n, _ := p.dtls.Read(buf[:])
      if n > 0 {
        log.Debug(n, " bytes of SCTP data received")
        p.sctp.Feed(buf[:n])
      }
    }
  }()

  handshakeDone := make(chan bool)
  go func () {
    var buf [1 << 16]byte
    tick := time.Tick(4 * time.Millisecond)
    for {
      select {
      case <-tick:
        n, _ := p.dtls.Spew(buf[:])
        if n > 0 {
          log.Debug(n, " bytes of DTLS data ready")
          rv, err := p.ice.Send(buf[:n])
          if err != nil {
            log.Warn(err)
          } else {
            log.Debug(rv, " bytes of DTLS data sent")
          }
        }
        continue
      case <-handshakeDone:
        close(handshakeDone)
        break
      }
      break
    }
  }()

  if err := p.dtls.Handshake(); err != nil {
    return err
  }
  handshakeDone <- true
  log.Debug("DTLS handshake done")

  go func () {
    var buf [1 << 16]byte
    for {
      data := <-p.sctp.BufferChannel
      log.Debug(len(data), " bytes of SCTP data ready")
      p.dtls.Write(data)

      n, _ := p.dtls.Spew(buf[:])
      if n > 0 {
        log.Debug(n, " bytes of DTLS data ready")
        rv, err := p.ice.Send(buf[:n])
        if err != nil {
          log.Warn(err)
        } else {
          log.Debug(rv, " bytes of DTLS data sent")
        }
      }
    }
  }()

  if p.role == roleClient {
    if err := p.sctp.Connect(p.remotePort); err != nil {
      return err
    }
  } else {
    if err := p.sctp.Accept(); err != nil {
      return err
    }
  }
  p.state = stateConnected
  log.Debug("SCTP handshake done")

  return nil
}

var numbers = []rune("0123456789")

func randSession() string {
  s := make([]rune, 16)
  rand.Seed(time.Now().UnixNano())
  for i := range s {
    s[i] = numbers[rand.Intn(10)]
  }
  return string(s)
}

func (p *Peer) GenerateOfferSdp() (string, error) {
  var offer []string
  offer = append(offer, "v=0")
  offer = append(offer, fmt.Sprintf("o=- %s 2 IN IP4 127.0.0.1", randSession()))
  offer = append(offer, "s=-")
  offer = append(offer, "t=0 0")
  offer = append(offer, "a=msid-semantic: WMS")
  offer = append(offer, "m=application 1 UDP/DTLS/SCTP webrtc-datachannel")
  offer = append(offer, "c=IN IP4 0.0.0.0")

  sdp := p.ice.GenerateSdp()
  sdps := strings.Split(sdp, "\n")
  for i := range sdps {
    if strings.HasPrefix(sdps[i], "a=ice-ufrag:") {
      offer = append(offer, sdps[i])
    } else if strings.HasPrefix(sdps[i], "a=ice-pwd:") {
      offer = append(offer, sdps[i])
    }
  }

  offer = append(offer, fmt.Sprintf("a=fingerprint:sha-256 %s", p.dtls.Fingerprint))
  if p.role == roleClient {
    offer = append(offer, "a=setup:active")
  } else {
    offer = append(offer, "a=setup:passive")
  }
  offer = append(offer, "a=mid:data")
  offer = append(offer, fmt.Sprintf("a=sctp-port:%d", p.sctp.Port))
  offer = append(offer, "")

  p.state = stateConnecting

  return strings.Join(offer, "\r\n"), nil
}

func (p *Peer) ParseOfferSdp(offer string) (int, error) {
  sdps := strings.Split(offer, "\r\n")
  for i := range sdps {
    if strings.HasPrefix(sdps[i], "a=sctp-port:") {
      sctpmap := strings.Split(sdps[i], " ")[0]
      port, err := strconv.Atoi(strings.Split(sctpmap, ":")[1])
      if err != nil {
        return 0, err
      }
      p.remotePort = port
    } else if strings.HasPrefix(sdps[i], "a=setup:active") {
      if p.role == roleClient {
        p.role = roleServer
      }
    } else if strings.HasPrefix(sdps[i], "a=setup:passive") {
      if p.role == roleServer {
        p.role = roleClient
      }
    }
  }

  rv, err := p.ice.ParseSdp(strings.Join(sdps, "\n"))
  if err != nil {
    return 0, err
  }

  p.state = stateConnecting

  return int(rv), nil
}

func (p *Peer) ParseCandidateSdp(cands string) int {
  sdps := strings.Split(cands, "\r\n")
  count := 0
  for i := range sdps {
    if sdps[i] == "" {
      continue
    }
    rv, err := p.ice.ParseCandidateSdp(sdps[i])
    if err != nil {
      continue
    }
    count = count + rv
  }
  return count
}
<|endoftext|>


<commit_before>package sketches

import (
	"strconv"
	"testing"

	"datamodel"
	pb "datamodel/protobuf"
	"utils"
)

func TestAddTopK(t *testing.T) {
	utils.SetupTests()
	defer utils.TearDownTests()

	info := datamodel.NewEmptyInfo()
	info.Properties.Size = utils.Int64p(3)
	info.Name = utils.Stringp("marvel")
	sketch, err := NewTopKSketch(info)

	if err != nil {
		t.Error("expected avengers to have no error, got", err)
	}

	values := [][]byte{
		[]byte("sabertooth"),
		[]byte("thunderbolt"),
		[]byte("cyclops"),
		[]byte("thunderbolt"),
		[]byte("thunderbolt"),
		[]byte("havoc"),
		[]byte("cyclops"),
		[]byte("cyclops"),
		[]byte("cyclops"),
		[]byte("havoc")}

	if _, err := sketch.Add(values); err != nil {
		t.Error("expected no errors, got", err)
	}

	type RankingsStruct struct {
		Value    string
		Position int64
		Count    int64
	}
	expectedRankings := make([]*RankingsStruct, 4, 4)

	expectedRankings[0] = &RankingsStruct{
		Value:    "cyclops",
		Count:    4,
		Position: 1,
	}

	expectedRankings[1] = &RankingsStruct{
		Value:    "thunderbolt",
		Count:    3,
		Position: 2,
	}

	expectedRankings[2] = &RankingsStruct{
		Value:    "havoc",
		Count:    2,
		Position: 3,
	}

	expectedRankings[3] = &RankingsStruct{
		Value:    "sabertooth",
		Count:    1,
		Position: 4,
	}

	if res, err := sketch.Get(values); err != nil {
		t.Error("expected no errors, got", err)
	} else {
		tmp := res.(*pb.RankingsResult)
		rres := tmp.GetRankings()
		for i := 0; i < len(rres); i++ {
			if expectedRankings[i].Value != rres[i].GetValue() {
				t.Errorf("expected ranking == %d, got %d", expectedRankings[i].Position, rres[i].GetValue())
			}
		}
	}
}

func BenchmarkTopK(b *testing.B) {
	values := make([][]byte, 10)
	for i := 0; i < 1024; i++ {
		avenger := "avenger" + strconv.Itoa(i)
		values = append(values, []byte(avenger))
	}

	for n := 0; n < b.N; n++ {
		info := datamodel.NewEmptyInfo()
		info.Properties.Size = utils.Int64p(1000)
		info.Name = utils.Stringp("marvel3")
		sketch, err := NewTopKSketch(info)
		if err != nil {
			b.Error("expected no errors, got", err)
		}
		for i := 0; i < 1000; i++ {
			if _, err := sketch.Add(values); err != nil {
				b.Error("expected no errors, got", err)
			}
		}
	}
}
<commit_msg>Bad Seif<commit_after>package sketches

import (
	"strconv"
	"testing"

	"datamodel"
	pb "datamodel/protobuf"
	"utils"
)

func TestAddTopK(t *testing.T) {
	utils.SetupTests()
	defer utils.TearDownTests()

	info := datamodel.NewEmptyInfo()
	info.Properties.Size = utils.Int64p(3)
	info.Name = utils.Stringp("marvel")
	sketch, err := NewTopKSketch(info)

	if err != nil {
		t.Error("expected avengers to have no error, got", err)
	}

	values := [][]byte{
		[]byte("sabertooth"),
		[]byte("thunderbolt"),
		[]byte("cyclops"),
		[]byte("thunderbolt"),
		[]byte("thunderbolt"),
		[]byte("havoc"),
		[]byte("cyclops"),
		[]byte("cyclops"),
		[]byte("cyclops"),
		[]byte("havoc")}

	if _, err := sketch.Add(values); err != nil {
		t.Error("expected no errors, got", err)
	}

	type RankingsStruct struct {
		Value    string
		Position int64
		Count    int64
	}
	expectedRankings := make([]*RankingsStruct, 4, 4)

	expectedRankings[0] = &RankingsStruct{
		Value:    "cyclops",
		Count:    4,
		Position: 1,
	}

	expectedRankings[1] = &RankingsStruct{
		Value:    "thunderbolt",
		Count:    3,
		Position: 2,
	}

	expectedRankings[2] = &RankingsStruct{
		Value:    "havoc",
		Count:    2,
		Position: 3,
	}

	expectedRankings[3] = &RankingsStruct{
		Value:    "sabertooth",
		Count:    1,
		Position: 4,
	}

	if res, err := sketch.Get(values); err != nil {
		t.Error("expected no errors, got", err)
	} else {
		tmp := res.(*pb.RankingsResult)
		rres := tmp.GetRankings()
		for i := 0; i < len(rres); i++ {
			if expectedRankings[i].Value != rres[i].GetValue() {
				t.Errorf("expected rank of %d := %s, got %s", expectedRankings[i].Position, expectedRankings[i].Value, rres[i].GetValue())
			}
		}
	}
}

func BenchmarkTopK(b *testing.B) {
	values := make([][]byte, 10)
	for i := 0; i < 1024; i++ {
		avenger := "avenger" + strconv.Itoa(i)
		values = append(values, []byte(avenger))
	}

	for n := 0; n < b.N; n++ {
		info := datamodel.NewEmptyInfo()
		info.Properties.Size = utils.Int64p(1000)
		info.Name = utils.Stringp("marvel3")
		sketch, err := NewTopKSketch(info)
		if err != nil {
			b.Error("expected no errors, got", err)
		}
		for i := 0; i < 1000; i++ {
			if _, err := sketch.Add(values); err != nil {
				b.Error("expected no errors, got", err)
			}
		}
	}
}
<|endoftext|>


<commit_before>package gorange

import (
	"context"
	"expvar"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"time"

	"github.com/karrick/gohm"
)

var errorCount = expvar.NewInt("errorCount")

// ProxyConfig specifies the configuration for a gorange proxy HTTP server.
type ProxyConfig struct {
	// CheckVersionPeriodicity directs the range proxy to periodically send the '%version' query
	// the proxied range servers, and if the value is greater than the previous value, to
	// asynchronously update all the values for recently used range keys.
	CheckVersionPeriodicity time.Duration

	// Log directs the proxy to emit common log formatted log lines to the specified io.Writer.
	Log io.Writer

	// LogFormat specifies the log format to use when Log is not nil.  See `gohm` package for
	// LogFormat specifications.  If left blank, uses `gohm.DefaultLogFormat`.
	LogFormat string

	// Port specifies which network port the proxy should bind to.
	Port uint

	// Servers specifies which addresses ought to be consulted as the source of truth for range
	// queries.
	Servers []string

	// Timeout specifies how long to wait for the source of truth to respond. If the zero-value,
	// no timeout will be used. Not having a timeout value may cause resource exhaustion where
	// any of the proxied servers take too long to return a response.
	Timeout time.Duration
}

// Proxy creates a proxy http server on the port that proxies range queries to the specified range
// servers.
func Proxy(config ProxyConfig) error {
	querier, err := NewQuerier(&Configurator{
		CheckVersionPeriodicity: config.CheckVersionPeriodicity,
		RetryCount:              len(config.Servers),
		Servers:                 config.Servers,
	})
	if err != nil {
		return err
	}

	mux := http.NewServeMux()
	mux.Handle("/range/expand", onlyGet(decodeURI(expand(querier))))
	mux.Handle("/range/list", onlyGet(decodeURI(list(querier))))
	mux.Handle("/", notFound())

	var h http.Handler = gohm.ConvertPanicsToErrors(mux)
	if config.Timeout > 0 {
		h = gohm.WithTimeout(config.Timeout, h)
	}
	if config.Log != nil {
		if config.LogFormat == "" {
			config.LogFormat = gohm.DefaultLogFormat
		}
		logBitmask := gohm.LogStatusAll
		h = gohm.LogStatusBitmaskWithFormat(config.LogFormat, &logBitmask, config.Log, h)
	}

	server := http.Server{
		Addr:         fmt.Sprintf(":%d", config.Port),
		Handler:      h,
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 15 * time.Second,
	}
	return server.ListenAndServe()
}

func notFound() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gohm.Error(w, r.URL.String(), http.StatusNotFound)
	})
}

func onlyGet(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "GET" {
			gohm.Error(w, r.Method, http.StatusMethodNotAllowed)
			return
		}
		next.ServeHTTP(w, r)
	})
}

type key int

const requestIDKey key = 0

func contextFromQuery(ctx context.Context, r *http.Request) (context.Context, error) {
	query, err := url.QueryUnescape(r.URL.RawQuery)
	if err != nil {
		return nil, fmt.Errorf("cannot decode query: %s", err)
	}
	return context.WithValue(ctx, requestIDKey, query), nil
}

func queryFromContext(ctx context.Context) string {
	return ctx.Value(requestIDKey).(string)
}

func decodeURI(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx, err := contextFromQuery(r.Context(), r)
		if err != nil {
			gohm.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

func expand(querier Querier) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		query := queryFromContext(r.Context())
		response, err := querier.Expand(query)
		if err != nil {
			gohm.Error(w, "cannot resolve query: "+err.Error(), http.StatusBadGateway)
			return
		}
		if _, err = io.WriteString(w, response); err != nil {
			gohm.Error(w, "cannot write response: "+err.Error(), http.StatusInternalServerError)
			return
		}
	})
}

func list(querier Querier) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		query := queryFromContext(r.Context())
		responses, err := querier.List(query)
		if err != nil {
			gohm.Error(w, "cannot resolve query: "+err.Error(), http.StatusBadGateway)
			return
		}
		for _, response := range responses {
			if _, err = fmt.Fprintf(w, "%s\r\n", response); err != nil {
				gohm.Error(w, "cannot write response: "+err.Error(), http.StatusInternalServerError)
				return
			}
		}
	})
}
<commit_msg>uses new gohm<commit_after>package gorange

import (
	"bytes"
	"context"
	"expvar"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"time"

	"github.com/karrick/gohm"
)

var errorCount = expvar.NewInt("errorCount")

// ProxyConfig specifies the configuration for a gorange proxy HTTP server.
type ProxyConfig struct {
	// CheckVersionPeriodicity directs the range proxy to periodically send the '%version' query
	// the proxied range servers, and if the value is greater than the previous value, to
	// asynchronously update all the values for recently used range keys.
	CheckVersionPeriodicity time.Duration

	// Log directs the proxy to emit common log formatted log lines to the specified io.Writer.
	Log io.Writer

	// LogFormat specifies the log format to use when Log is not nil.  See `gohm` package for
	// LogFormat specifications.  If left blank, uses `gohm.DefaultLogFormat`.
	LogFormat string

	// Port specifies which network port the proxy should bind to.
	Port uint

	// Servers specifies which addresses ought to be consulted as the source of truth for range
	// queries.
	Servers []string

	// Timeout specifies how long to wait for the source of truth to respond. If the zero-value,
	// no timeout will be used. Not having a timeout value may cause resource exhaustion where
	// any of the proxied servers take too long to return a response.
	Timeout time.Duration
}

// Proxy creates a proxy http server on the port that proxies range queries to the specified range
// servers.
func Proxy(config ProxyConfig) error {
	querier, err := NewQuerier(&Configurator{
		CheckVersionPeriodicity: config.CheckVersionPeriodicity,
		RetryCount:              len(config.Servers),
		Servers:                 config.Servers,
	})
	if err != nil {
		return err
	}

	mux := http.NewServeMux()
	mux.Handle("/range/expand", onlyGet(decodeURI(expand(querier))))
	mux.Handle("/range/list", onlyGet(decodeURI(list(querier))))
	mux.Handle("/", notFound()) // while not required, this makes for a nicer log output and client response

	var h http.Handler = gohm.New(mux, gohm.Config{
		LogWriter: config.Log,
		Timeout:   config.Timeout,
	})

	server := http.Server{
		Addr:         fmt.Sprintf(":%d", config.Port),
		Handler:      h,
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 15 * time.Second,
	}
	return server.ListenAndServe()
}

func notFound() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gohm.Error(w, r.URL.String(), http.StatusNotFound)
	})
}

func onlyGet(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "GET" {
			gohm.Error(w, r.Method, http.StatusMethodNotAllowed)
			return
		}
		next.ServeHTTP(w, r)
	})
}

type key int

const requestIDKey key = 0

func contextFromQuery(ctx context.Context, r *http.Request) (context.Context, error) {
	query, err := url.QueryUnescape(r.URL.RawQuery)
	if err != nil {
		return nil, fmt.Errorf("cannot decode query: %s", err)
	}
	return context.WithValue(ctx, requestIDKey, query), nil
}

func queryFromContext(ctx context.Context) string {
	return ctx.Value(requestIDKey).(string)
}

func decodeURI(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx, err := contextFromQuery(r.Context(), r)
		if err != nil {
			gohm.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

func expand(querier Querier) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		query := queryFromContext(r.Context())
		response, err := querier.Expand(query)
		if err != nil {
			gohm.Error(w, "cannot resolve query: "+err.Error(), http.StatusBadGateway)
			return
		}
		if _, err = io.WriteString(w, response); err != nil {
			gohm.Error(w, "cannot write response: "+err.Error(), http.StatusInternalServerError)
			return
		}
	})
}

func list(querier Querier) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		query := queryFromContext(r.Context())
		responses, err := querier.List(query)
		if err != nil {
			gohm.Error(w, "cannot resolve query: "+err.Error(), http.StatusBadGateway)
			return
		}
		buf := new(bytes.Buffer)
		for _, response := range responses {
			if _, err = buf.WriteString(response + "\r\n"); err != nil {
				gohm.Error(w, "cannot write response: "+err.Error(), http.StatusInternalServerError)
				return
			}
		}
		if _, err = buf.WriteTo(w); err != nil {
			gohm.Error(w, "cannot write response: "+err.Error(), http.StatusInternalServerError)
			return
		}
	})
}
<|endoftext|>