<commit_before>package main

// this is a direct port of arc90's readbility javascript
// http://code.google.com/p/arc90labs-readability

import (
	"code.google.com/p/cascadia"
	"code.google.com/p/go.net/html"
	"code.google.com/p/go.net/html/atom"
	"fmt"
	//	"github.com/matrixik/goquery"
	"regexp"
	"math"
	//	"sort"
	"strings"

	"os"
)

type CandidateMap map[*html.Node]*Candidate

func (candidates CandidateMap) get(n *html.Node) *Candidate {
	c, ok := candidates[n]
	if !ok {
		c = newCandidate(n, "")
		candidates[n] = c
	}
	return c
}

func (cm *CandidateMap) initializeNode(node *html.Node) {
	c := cm.get(node)

	switch node.DataAtom {
	case atom.Div:
		c.addScore(5, "<div>")
	case atom.Pre, atom.Td, atom.Blockquote:
		c.addScore(3, "<pre>, <td> or <blockquote>")
	case atom.Address, atom.Ol, atom.Ul, atom.Dl, atom.Dd, atom.Li, atom.Form:
		c.addScore(-3, "address, list or form")
	case atom.H1, atom.H2, atom.H3, atom.H4, atom.H5, atom.H6, atom.Th:
		c.addScore(-5, "heading")
	}

	if score := getClassWeight(node); score != 0 {
		c.addScore(score, "class/id score")
	}
}

func DoVoodoo(root *html.Node) {
	removeScripts(root)
	// TODO: Turn all double br's into p's? Kill <style> tags? (see prepDocument())
	grabArticle(root)
}

// remove all <script> elements
func removeScripts(root *html.Node) {
	sel := cascadia.MustCompile("script")
	for _, script := range sel.MatchAll(root) {
		script.Parent.RemoveChild(script)
	}
}

var unlikelyCandidates = regexp.MustCompile(`(?i)combx|comment|community|disqus|extra|foot|header|menu|remark|rss|shoutbox|sidebar|sponsor|ad-break|agegate|pagination|pager|popup|tweet|twitter`)
var okMaybeItsACandidate = regexp.MustCompile(`(?i)and|article|body|column|main|shadow`)

var positivePat = regexp.MustCompile(`(?i)article|body|content|entry|hentry|main|page|pagination|post|text|blog|story`)
var negativePat = regexp.MustCompile(`(?i)combx|comment|com-|contact|foot|footer|footnote|masthead|media|meta|outbrain|promo|related|scroll|shoutbox|sidebar|sponsor|shopping|tags|tool|widget`)

func grabArticle(root *html.Node) {

	var candidates = make(CandidateMap)

	stripUnlikelyCandidates := true

	/**
	 * First, node prepping. Trash nodes that look cruddy (like ones with the class name "comment", etc), and turn divs
	 * into P tags where they have been used inappropriately (as in, where they contain no other block level elements.)
	 *
	 * Note: Assignment from index for performance. See http://www.peachpit.com/articles/article.aspx?p=31567&seqNum=5
	 * TODO: Shouldn't this be a reverse traversal?
	**/

	nodesToScore := make([]*html.Node, 0, 128)

	allNodes := cascadia.MustCompile("*")
	for _, node := range allNodes.MatchAll(root) {
		if stripUnlikelyCandidates {
			unlikelyMatchString := getAttr(node, "class") + getAttr(node, "id")
			//		fmt.Printf("? %s\n",unlikelyMatchString)
			if unlikelyCandidates.MatchString(unlikelyMatchString) == true &&
				okMaybeItsACandidate.MatchString(unlikelyMatchString) == false &&
				node.DataAtom != atom.Body {
				fmt.Printf("Removing unlikely candidate - %s\n", describeNode(node))
				node.Parent.RemoveChild(node)
				continue
			}
		}

		if node.DataAtom == atom.P || node.DataAtom == atom.Td || node.DataAtom == atom.Pre {
			nodesToScore = append(nodesToScore, node)
		}
		/* XYZZY TODO: Turn all divs that don't have children block level elements into p's */
	}

	fmt.Printf("%d nodes to score\n", len(nodesToScore))

	/*
	 * Loop through all paragraphs, and assign a score to them based on how content-y they look.
	 * Then add their score to their parent node.
	 *
	 * A score is determined by things like number of commas, class names, etc. Maybe eventually link density.
	 */

	// TODO: make sure order is right, so proper scores propogate
	for _, node := range nodesToScore {
		parentNode := node.Parent
		var grandParentNode *html.Node = nil
		if parentNode != nil {
			grandParentNode = parentNode.Parent
		}

		innerText := getTextContent(node)
		/* If this paragraph is less than 25 characters, don't even count it. */
		if len(innerText) < 25 {
			continue
		}

		if _, exists := candidates[parentNode]; !exists {
			candidates.initializeNode(parentNode)
		}
		if grandParentNode != nil {
			if _, exists := candidates[grandParentNode]; !exists {
				candidates.initializeNode(grandParentNode)
			}
		}

		contentScore := 1.0

		// add points for any commas
		contentScore += float64(strings.Count(innerText, ","))

		// 1 point for every 100 bytes in this para, up to 3 points
		foo := float64(len(innerText)) / 100
		if foo > 3 {
			foo = 3
		}
		contentScore += foo

		/* Add the content score to the parent. The grandparent gets half. */
		candidates.get(parentNode).addScore(contentScore, "Child content")
		if grandParentNode != nil {
			halfScore := contentScore / 2
			if halfScore > 0 {
				candidates.get(grandParentNode).addScore(halfScore, "Child content")
			}
		}
	}

	/**
	 * Scale the final candidates score based on link density. Good content should have a
	 * relatively small link density (5% or less) and be mostly unaffected by this operation.
	 **/
	for _, c := range candidates {
		c.scaleScore((1 - getLinkDensity(c.Node)), "link density")
	}

	/**
	 * After we've calculated scores, loop through all of the possible candidate nodes we found
	 * and find the one with the highest score.
	**/
	var topCandidate *Candidate = nil
	for _, c := range candidates {
		if topCandidate == nil || c.TotalScore > topCandidate.TotalScore {
			topCandidate = c
		}
	}

	/**
	 * Now that we have the top candidate, look through its siblings for content that might also be related.
	 * Things like preambles, content split by ads that we removed, etc.
	**/

	siblingScoreThreshold := topCandidate.TotalScore * 0.2
	if siblingScoreThreshold < 10 {
		siblingScoreThreshold = 10
	}

	contentNodes := make([]*html.Node, 0, 64)

	for siblingNode := topCandidate.Node.Parent.FirstChild; siblingNode != nil; siblingNode = siblingNode.NextSibling {
		useIt := false
		if siblingNode == topCandidate.Node {
			useIt = true
		} else {

			contentBonus := 0.0
			/* Give a bonus if sibling nodes and top candidates have the exact same classname */
			topClass := getAttr(topCandidate.Node, "class")
			if getAttr(siblingNode, "class") == topClass && topClass != "" {
				contentBonus += topCandidate.TotalScore * 0.2
			}
			if sc, ok := candidates[siblingNode]; ok == true && sc.TotalScore+contentBonus >= siblingScoreThreshold {
				useIt = true
			}

			if siblingNode.DataAtom == atom.P {
				linkDensity := getLinkDensity(siblingNode)
				nodeContent := getTextContent(siblingNode)
				nodeLength := len(nodeContent)

				if nodeLength >= 80 && linkDensity < 0.25 {
					useIt = true
				} else if nodeLength < 80 && linkDensity == 0 && regexp.MustCompile(`\.( |$)`).MatchString(nodeContent) {
					useIt = true
				}
			}
		}

		if useIt {
			contentNodes = append(contentNodes, siblingNode)
		}

	}

	//	for _, candidate := range candidates {
	//		candidate.dump()
	//	}
	//	fmt.Printf("best:\n")
	//	topCandidate.dump()

	// go through and clean out any cruft
	for _,node := range contentNodes {
        cleanConditionally(node, "form",candidates)
        cleanConditionally(node, "table",candidates)
        cleanConditionally(node, "ul",candidates)
        cleanConditionally(node, "div",candidates)
	}

	fmt.Printf("picked %d nodes:\n", len(contentNodes))
	for _, n := range contentNodes {
		fmt.Printf("%s:\n", describeNode(n))
		html.Render(os.Stdout,n)
	}

}

/*
 * Get an elements class/id weight. Uses regular expressions to tell if this
 * element looks good or bad.
**/
func getClassWeight(n *html.Node) float64 {
	//if(!readability.flagIsActive(readability.FLAG_WEIGHT_CLASSES)) {
	//    return 0;
	//}

	score := 0.0

	cls := getAttr(n, "class")
	id := getAttr(n, "id")

	/* Look for a special classname */
	if negativePat.MatchString(cls) {
		score -= 25
	}
	if positivePat.MatchString(cls) {
		score += 25
	}
	/* Look for a special ID */
	if negativePat.MatchString(id) {
		score -= 25
	}
	if positivePat.MatchString(id) {
		score += 25
	}

	return score
}

/*
 * Prepare the article nodes for display. Clean out any inline styles,
 * iframes, forms, strip extraneous <p> tags, etc.
 *
 */
func prepArticle(articleContent []*html.Node) {
}

/**
 * Clean an element of all tags of type "tag" if they look fishy.
 * "Fishy" is an algorithm based on content length, classnames, link density, number of images & embeds, etc.
 **/
func cleanConditionally(e *html.Node, tagSel string, candidates CandidateMap) {
	sel := cascadia.MustCompile(tagSel)
	toRemove := false
	doomed := make([]*html.Node,0,32)
	for _, node := range sel.MatchAll(e) {
		weight := getClassWeight(node)
		var contentScore float64 = 0.0
		if c, ok := candidates[node]; ok {
			contentScore = c.TotalScore
		}

		if weight+contentScore < 0 {
			toRemove = true
		} else {
			textContent := getTextContent(node)
			if strings.Count(textContent, ",") < 10 {
				/*
				 * If there are not very many commas, and the number of
				 * non-paragraph elements is more than paragraphs or other ominous signs, remove the element.
				 */
				p := len(cascadia.MustCompile("p").MatchAll(node))
				img := len(cascadia.MustCompile("img").MatchAll(node))
				li := len(cascadia.MustCompile("li").MatchAll(node))
				input := len(cascadia.MustCompile("input").MatchAll(node))
				// XYZZY TODO: exclude videos?
				embedCount := len(cascadia.MustCompile("embed").MatchAll(node))
				linkDensity := getLinkDensity(node)

				if  img > p  {
                    toRemove = true
                } else if li > p && node.DataAtom != atom.Ul && node.DataAtom != atom.Ol {
                    toRemove = true
                } else if input > int(math.Floor(float64(p)/3.0)) {
                    toRemove = true
                } else if len(textContent) < 25 && (img == 0 || img > 2) {
                    toRemove = true
                } else if weight < 25 && linkDensity > 0.2 {
                    toRemove = true
                } else if weight >= 25 && linkDensity > 0.5 {
                    toRemove = true
                } else if (embedCount == 1 && len(textContent) < 75) || embedCount > 1 {
                    toRemove = true
                }

			}
		}

		if toRemove {
			doomed = append(doomed,node)
		}
	}

	for _,n := range(doomed) {
		if n.Parent != nil {
			n.Parent.RemoveChild(n)
		}
	}

}
<commit_msg>add some more content-cleaning code<commit_after>package main

// this is a direct port of arc90's readbility javascript
// http://code.google.com/p/arc90labs-readability

import (
	"code.google.com/p/cascadia"
	"code.google.com/p/go.net/html"
	"code.google.com/p/go.net/html/atom"
	"fmt"
	//	"github.com/matrixik/goquery"
	"math"
	"regexp"
	//	"sort"
	"strings"

	"os"
)

type CandidateMap map[*html.Node]*Candidate

func (candidates CandidateMap) get(n *html.Node) *Candidate {
	c, ok := candidates[n]
	if !ok {
		c = newCandidate(n, "")
		candidates[n] = c
	}
	return c
}

func (cm *CandidateMap) initializeNode(node *html.Node) {
	c := cm.get(node)

	switch node.DataAtom {
	case atom.Div:
		c.addScore(5, "<div>")
	case atom.Pre, atom.Td, atom.Blockquote:
		c.addScore(3, "<pre>, <td> or <blockquote>")
	case atom.Address, atom.Ol, atom.Ul, atom.Dl, atom.Dd, atom.Li, atom.Form:
		c.addScore(-3, "address, list or form")
	case atom.H1, atom.H2, atom.H3, atom.H4, atom.H5, atom.H6, atom.Th:
		c.addScore(-5, "heading")
	}

	if score := getClassWeight(node); score != 0 {
		c.addScore(score, "class/id score")
	}
}

func DoVoodoo(root *html.Node) {
	removeScripts(root)
	// TODO: Turn all double br's into p's? Kill <style> tags? (see prepDocument())
	grabArticle(root)
}

// remove all <script> elements
func removeScripts(root *html.Node) {
	sel := cascadia.MustCompile("script")
	for _, script := range sel.MatchAll(root) {
		script.Parent.RemoveChild(script)
	}
}

var unlikelyCandidates = regexp.MustCompile(`(?i)combx|comment|community|disqus|extra|foot|header|menu|remark|rss|shoutbox|sidebar|sponsor|ad-break|agegate|pagination|pager|popup|tweet|twitter`)
var okMaybeItsACandidate = regexp.MustCompile(`(?i)and|article|body|column|main|shadow`)

var positivePat = regexp.MustCompile(`(?i)article|body|content|entry|hentry|main|page|pagination|post|text|blog|story`)
var negativePat = regexp.MustCompile(`(?i)combx|comment|com-|contact|foot|footer|footnote|masthead|media|meta|outbrain|promo|related|scroll|shoutbox|sidebar|sponsor|shopping|tags|tool|widget`)

func grabArticle(root *html.Node) {

	var candidates = make(CandidateMap)

	stripUnlikelyCandidates := true

	/**
	 * First, node prepping. Trash nodes that look cruddy (like ones with the class name "comment", etc), and turn divs
	 * into P tags where they have been used inappropriately (as in, where they contain no other block level elements.)
	 *
	 * Note: Assignment from index for performance. See http://www.peachpit.com/articles/article.aspx?p=31567&seqNum=5
	 * TODO: Shouldn't this be a reverse traversal?
	**/

	nodesToScore := make([]*html.Node, 0, 128)

	allNodes := cascadia.MustCompile("*")
	for _, node := range allNodes.MatchAll(root) {
		if stripUnlikelyCandidates {
			unlikelyMatchString := getAttr(node, "class") + getAttr(node, "id")
			//		fmt.Printf("? %s\n",unlikelyMatchString)
			if unlikelyCandidates.MatchString(unlikelyMatchString) == true &&
				okMaybeItsACandidate.MatchString(unlikelyMatchString) == false &&
				node.DataAtom != atom.Body {
				fmt.Printf("Removing unlikely candidate - %s\n", describeNode(node))
				node.Parent.RemoveChild(node)
				continue
			}
		}

		if node.DataAtom == atom.P || node.DataAtom == atom.Td || node.DataAtom == atom.Pre {
			nodesToScore = append(nodesToScore, node)
		}
		/* XYZZY TODO: Turn all divs that don't have children block level elements into p's */
	}

	fmt.Printf("%d nodes to score\n", len(nodesToScore))

	/*
	 * Loop through all paragraphs, and assign a score to them based on how content-y they look.
	 * Then add their score to their parent node.
	 *
	 * A score is determined by things like number of commas, class names, etc. Maybe eventually link density.
	 */

	// TODO: make sure order is right, so proper scores propogate
	for _, node := range nodesToScore {
		parentNode := node.Parent
		var grandParentNode *html.Node = nil
		if parentNode != nil {
			grandParentNode = parentNode.Parent
		}

		innerText := getTextContent(node)
		/* If this paragraph is less than 25 characters, don't even count it. */
		if len(innerText) < 25 {
			continue
		}

		if _, exists := candidates[parentNode]; !exists {
			candidates.initializeNode(parentNode)
		}
		if grandParentNode != nil {
			if _, exists := candidates[grandParentNode]; !exists {
				candidates.initializeNode(grandParentNode)
			}
		}

		contentScore := 1.0

		// add points for any commas
		contentScore += float64(strings.Count(innerText, ","))

		// 1 point for every 100 bytes in this para, up to 3 points
		foo := float64(len(innerText)) / 100
		if foo > 3 {
			foo = 3
		}
		contentScore += foo

		/* Add the content score to the parent. The grandparent gets half. */
		candidates.get(parentNode).addScore(contentScore, "Child content")
		if grandParentNode != nil {
			halfScore := contentScore / 2
			if halfScore > 0 {
				candidates.get(grandParentNode).addScore(halfScore, "Child content")
			}
		}
	}

	/**
	 * Scale the final candidates score based on link density. Good content should have a
	 * relatively small link density (5% or less) and be mostly unaffected by this operation.
	 **/
	for _, c := range candidates {
		c.scaleScore((1 - getLinkDensity(c.Node)), "link density")
	}

	/**
	 * After we've calculated scores, loop through all of the possible candidate nodes we found
	 * and find the one with the highest score.
	**/
	var topCandidate *Candidate = nil
	for _, c := range candidates {
		if topCandidate == nil || c.TotalScore > topCandidate.TotalScore {
			topCandidate = c
		}
	}

	/**
	 * Now that we have the top candidate, look through its siblings for content that might also be related.
	 * Things like preambles, content split by ads that we removed, etc.
	**/

	siblingScoreThreshold := topCandidate.TotalScore * 0.2
	if siblingScoreThreshold < 10 {
		siblingScoreThreshold = 10
	}

	contentNodes := make([]*html.Node, 0, 64)

	for siblingNode := topCandidate.Node.Parent.FirstChild; siblingNode != nil; siblingNode = siblingNode.NextSibling {
		useIt := false
		if siblingNode == topCandidate.Node {
			useIt = true
		} else {

			contentBonus := 0.0
			/* Give a bonus if sibling nodes and top candidates have the exact same classname */
			topClass := getAttr(topCandidate.Node, "class")
			if getAttr(siblingNode, "class") == topClass && topClass != "" {
				contentBonus += topCandidate.TotalScore * 0.2
			}
			if sc, ok := candidates[siblingNode]; ok == true && sc.TotalScore+contentBonus >= siblingScoreThreshold {
				useIt = true
			}

			if siblingNode.DataAtom == atom.P {
				linkDensity := getLinkDensity(siblingNode)
				nodeContent := getTextContent(siblingNode)
				nodeLength := len(nodeContent)

				if nodeLength >= 80 && linkDensity < 0.25 {
					useIt = true
				} else if nodeLength < 80 && linkDensity == 0 && regexp.MustCompile(`\.( |$)`).MatchString(nodeContent) {
					useIt = true
				}
			}
		}

		if useIt {
			contentNodes = append(contentNodes, siblingNode)
		}

	}

	removeCruft(contentNodes, candidates)
	sanitiseContent(contentNodes, candidates)

	fmt.Printf("extracted %d nodes:\n", len(contentNodes))
	for _, n := range contentNodes {
		fmt.Printf("%s:\n", describeNode(n))
		html.Render(os.Stdout, n)
	}

}

/*
 * Get an elements class/id weight. Uses regular expressions to tell if this
 * element looks good or bad.
**/
func getClassWeight(n *html.Node) float64 {
	//if(!readability.flagIsActive(readability.FLAG_WEIGHT_CLASSES)) {
	//    return 0;
	//}

	score := 0.0

	cls := getAttr(n, "class")
	id := getAttr(n, "id")

	/* Look for a special classname */
	if negativePat.MatchString(cls) {
		score -= 25
	}
	if positivePat.MatchString(cls) {
		score += 25
	}
	/* Look for a special ID */
	if negativePat.MatchString(id) {
		score -= 25
	}
	if positivePat.MatchString(id) {
		score += 25
	}

	return score
}

// Remove all extraneous crap in the content - related articles, share buttons etc...
// (equivalent to prepArticle() in readbility.js)
func removeCruft(contentNodes []*html.Node, candidates CandidateMap) {

	zapConditionally(contentNodes, "form", candidates)
	zap(contentNodes, "object")
	zap(contentNodes, "h1")

	// If there is only one h2, they are probably using it
	// as a header and not a subheader, so remove it since we already have a header.
	h2Count := 0
	h2Sel := cascadia.MustCompile("h2")
	for _, node := range contentNodes {
		h2Count += len(h2Sel.MatchAll(node))
	}

	if h2Count == 1 {
		zap(contentNodes, "h2")
	}
	zap(contentNodes, "iframe")

	//cleanHeaders()

	/* Do these last as the previous stuff may have removed junk that will affect these */
	zapConditionally(contentNodes, "table", candidates)
	zapConditionally(contentNodes, "ul", candidates)
	zapConditionally(contentNodes, "div", candidates)
}

func zap(contentNodes []*html.Node, tagSel string) {
	doomed := make([]*html.Node, 0, 32)
	sel := cascadia.MustCompile(tagSel)
	for _, contentNode := range contentNodes {
		for _, node := range sel.MatchAll(contentNode) {
			// XYZZY TODO: preserve videos?
			doomed = append(doomed, node)
		}
	}
	for _, n := range doomed {
		if n.Parent != nil {
			n.Parent.RemoveChild(n)
		}
	}
}

/**
 * Clean a set of elements, removing all matching tags if they look fishy.
 * "Fishy" is an algorithm based on content length, classnames, link density, number of images & embeds, etc.
 **/
func zapConditionally(contentNodes []*html.Node, tagSel string, candidates CandidateMap) {

	doomed := make([]*html.Node, 0, 32)
	sel := cascadia.MustCompile(tagSel)
	for _, e := range contentNodes {
		toRemove := false
		for _, node := range sel.MatchAll(e) {
			weight := getClassWeight(node)
			var contentScore float64 = 0.0
			if c, ok := candidates[node]; ok {
				contentScore = c.TotalScore
			}

			if weight+contentScore < 0 {
				toRemove = true
			} else {
				textContent := getTextContent(node)
				if strings.Count(textContent, ",") < 10 {
					/*
					 * If there are not very many commas, and the number of
					 * non-paragraph elements is more than paragraphs or other ominous signs, remove the element.
					 */
					p := len(cascadia.MustCompile("p").MatchAll(node))
					img := len(cascadia.MustCompile("img").MatchAll(node))
					li := len(cascadia.MustCompile("li").MatchAll(node))
					input := len(cascadia.MustCompile("input").MatchAll(node))
					// XYZZY TODO: exclude videos?
					embedCount := len(cascadia.MustCompile("embed").MatchAll(node))
					linkDensity := getLinkDensity(node)

					if img > p {
						toRemove = true
					} else if li > p && node.DataAtom != atom.Ul && node.DataAtom != atom.Ol {
						toRemove = true
					} else if input > int(math.Floor(float64(p)/3.0)) {
						toRemove = true
					} else if len(textContent) < 25 && (img == 0 || img > 2) {
						toRemove = true
					} else if weight < 25 && linkDensity > 0.2 {
						toRemove = true
					} else if weight >= 25 && linkDensity > 0.5 {
						toRemove = true
					} else if (embedCount == 1 && len(textContent) < 75) || embedCount > 1 {
						toRemove = true
					}

				}
			}

			if toRemove {
				doomed = append(doomed, node)
			}
		}
	}

	for _, n := range doomed {
		if n.Parent != nil {
			n.Parent.RemoveChild(n)
		}
	}
}

func filterAttrs(n *html.Node, fn func(*html.Attribute) bool) {
	var out []html.Attribute // == nil
	for _, a := range n.Attr {
		if fn(&a) {
			out = append(out, a)
		}
	}
	n.Attr = out
}

// Tidy up extracted content into something that'll produce reasonable html when
// rendered
func sanitiseContent(contentNodes []*html.Node, candidates CandidateMap) {

	var commentSel cascadia.Selector = func(n *html.Node) bool {
		return n.Type == html.CommentNode
	}
	var textSel cascadia.Selector = func(n *html.Node) bool {
		return n.Type == html.TextNode
	}
	var elementSel cascadia.Selector = func(n *html.Node) bool {
		return n.Type == html.ElementNode
	}

	for _, node := range contentNodes {
		// remove all comments
		for _, n := range commentSel.MatchAll(node) {
			n.Parent.RemoveChild(n)
		}
		// trim leading/trailing space in text
		for _, n := range textSel.MatchAll(node) {
			n.Data = strings.TrimSpace(n.Data)
		}
		// remove styles, ids and classes
		for _, n := range elementSel.MatchAll(node) {
			filterAttrs(n, func(attr *html.Attribute) bool {
				return attr.Key != "style" && attr.Key != "id" && attr.Key != "class"
			})
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"strings"

	"github.com/flynn/flynn-host/client"
	"github.com/flynn/flynn-host/sampi"
	"github.com/flynn/flynn-host/types"
	"github.com/flynn/go-discoverd"
	"github.com/flynn/go-dockerclient"
	rpc "github.com/flynn/rpcplus/comborpc"
	"github.com/technoweenie/grohl"
)

func main() {
	hostname, _ := os.Hostname()
	externalAddr := flag.String("external", "", "external IP of host")
	configFile := flag.String("config", "", "configuration file")
	manifestFile := flag.String("manifest", "", "manifest file")
	hostID := flag.String("id", hostname, "host id")
	flag.Parse()
	grohl.AddContext("app", "lorne")
	grohl.Log(grohl.Data{"at": "start"})
	g := grohl.NewContext(grohl.Data{"fn": "main"})

	dockerc, err := docker.NewClient("unix:///var/run/docker.sock")
	if err != nil {
		log.Fatal(err)
	}

	state := NewState()
	ports := make(chan int)

	go allocatePorts(ports, 55000, 65535)
	go serveHTTP(&Host{state: state, docker: dockerc}, &attachHandler{state: state, docker: dockerc})
	go streamEvents(dockerc, state)

	processor := &jobProcessor{
		externalAddr: *externalAddr,
		docker:       dockerc,
		state:        state,
		discoverd:    os.Getenv("DISCOVERD"),
	}

	runner := &manifestRunner{
		env:        parseEnviron(),
		externalIP: *externalAddr,
		ports:      ports,
		processor:  processor,
		docker:     dockerc,
	}

	var disc *discoverd.Client
	if *manifestFile != "" {
		f, err := os.Open(*manifestFile)
		if err != nil {
			log.Fatal(err)
		}
		services, err := runner.runManifest(f)
		if err != nil {
			log.Fatal(err)
		}
		f.Close()

		if d, ok := services["discoverd"]; ok {
			processor.discoverd = fmt.Sprintf("%s:%d", d.InternalIP, d.TCPPorts[0])
			disc, err = discoverd.NewClientUsingAddress(processor.discoverd)
			if err != nil {
				log.Fatal(err)
			}
		}
	}

	if processor.discoverd == "" && *externalAddr != "" {
		processor.discoverd = *externalAddr + ":1111"
	}
	// HACK: use env as global for discoverd connection in sampic
	os.Setenv("DISCOVERD", processor.discoverd)
	if disc == nil {
		disc, err = discoverd.NewClientUsingAddress(processor.discoverd)
		if err != nil {
			log.Fatal(err)
		}
	}
	sampiStandby, err := disc.RegisterAndStandby("flynn-host", *externalAddr+":1113", map[string]string{"id": *hostID})
	if err != nil {
		log.Fatal(err)
	}
	go func() {
		<-sampiStandby
		rpc.Register(sampi.NewCluster(sampi.NewState()))
	}()

	cluster, err := client.New()
	if err != nil {
		log.Fatal(err)
	}
	g.Log(grohl.Data{"at": "sampi_connected"})

	events := make(chan host.Event)
	state.AddListener("all", events)
	go syncScheduler(cluster, events)

	var h *host.Host
	if *configFile != "" {
		h, err = openConfig(*configFile)
		if err != nil {
			log.Fatal(err)
		}
	} else {
		h = &host.Host{Resources: make(map[string]host.ResourceValue)}
	}
	if _, ok := h.Resources["memory"]; !ok {
		h.Resources["memory"] = host.ResourceValue{Value: 1024}
	}
	h.ID = *hostID
	h.Jobs = state.ClusterJobs()

	jobs := make(chan *host.Job)
	hostErr := cluster.ConnectHost(h, jobs)
	g.Log(grohl.Data{"at": "host_registered"})
	processor.Process(ports, jobs)
	log.Fatal(*hostErr)
}

type jobProcessor struct {
	externalAddr string
	discoverd    string
	docker       interface {
		CreateContainer(*docker.Config) (*docker.Container, error)
		PullImage(docker.PullImageOptions, io.Writer) error
		StartContainer(string, *docker.HostConfig) error
	}
	state *State
}

func (p *jobProcessor) Process(ports <-chan int, jobs chan *host.Job) {
	for job := range jobs {
		p.processJob(ports, job)
	}
}

func (p *jobProcessor) processJob(ports <-chan int, job *host.Job) (*docker.Container, error) {
	g := grohl.NewContext(grohl.Data{"fn": "process_job", "job.id": job.ID})
	g.Log(grohl.Data{"at": "start", "job.image": job.Config.Image, "job.cmd": job.Config.Cmd, "job.entrypoint": job.Config.Entrypoint})

	for i := 0; i < job.TCPPorts; i++ {
		port := strconv.Itoa(<-ports)
		if i == 0 {
			job.Config.Env = append(job.Config.Env, "PORT="+port)
		}
		job.Config.Env = append(job.Config.Env, fmt.Sprintf("PORT_%d=%s", i, port))
		if job.Config.ExposedPorts == nil {
			job.Config.ExposedPorts = make(map[string]struct{}, job.TCPPorts)
		}
		job.Config.ExposedPorts[port+"/tcp"] = struct{}{}
		if job.HostConfig == nil {
			job.HostConfig = &docker.HostConfig{
				PortBindings:    make(map[string][]docker.PortBinding, job.TCPPorts),
				PublishAllPorts: true,
			}
		}
		job.HostConfig.PortBindings[port+"/tcp"] = []docker.PortBinding{{HostPort: port}}
	}
	if p.externalAddr != "" {
		job.Config.Env = appendUnique(job.Config.Env, "EXTERNAL_IP="+p.externalAddr, "SD_HOST="+p.externalAddr, "DISCOVERD="+p.discoverd)
	}
	p.state.AddJob(job)
	g.Log(grohl.Data{"at": "create_container"})
	container, err := p.docker.CreateContainer(job.Config)
	if err == docker.ErrNoSuchImage {
		g.Log(grohl.Data{"at": "pull_image"})
		err = p.docker.PullImage(docker.PullImageOptions{Repository: job.Config.Image}, os.Stdout)
		if err != nil {
			g.Log(grohl.Data{"at": "pull_image", "status": "error", "err": err})
			p.state.SetStatusFailed(job.ID, err)
			return nil, err
		}
		container, err = p.docker.CreateContainer(job.Config)
	}
	if err != nil {
		g.Log(grohl.Data{"at": "create_container", "status": "error", "err": err})
		p.state.SetStatusFailed(job.ID, err)
		return nil, err
	}
	p.state.SetContainerID(job.ID, container.ID)
	p.state.WaitAttach(job.ID)
	g.Log(grohl.Data{"at": "start_container"})
	if err := p.docker.StartContainer(container.ID, job.HostConfig); err != nil {
		g.Log(grohl.Data{"at": "start_container", "status": "error", "err": err})
		p.state.SetStatusFailed(job.ID, err)
		return nil, err
	}
	p.state.SetStatusRunning(job.ID)
	g.Log(grohl.Data{"at": "finish"})
	return container, nil
}

func appendUnique(s []string, vars ...string) []string {
outer:
	for _, v := range vars {
		for _, existing := range s {
			if strings.HasPrefix(existing, strings.SplitN(v, "=", 2)[0]+"=") {
				continue outer
			}
		}
		s = append(s, v)
	}
	return s
}

type sampiSyncClient interface {
	RemoveJobs([]string) error
}

func syncScheduler(scheduler sampiSyncClient, events <-chan host.Event) {
	for event := range events {
		if event.Event != "stop" {
			continue
		}
		grohl.Log(grohl.Data{"fn": "scheduler_event", "at": "remove_job", "job.id": event.JobID})
		if err := scheduler.RemoveJobs([]string{event.JobID}); err != nil {
			grohl.Log(grohl.Data{"fn": "scheduler_event", "at": "remove_job", "status": "error", "err": err, "job.id": event.JobID})
			// TODO: try to reconnect?
		}
	}
}

type dockerStreamClient interface {
	Events() (*docker.EventStream, error)
	InspectContainer(string) (*docker.Container, error)
}

func streamEvents(client dockerStreamClient, state *State) {
	stream, err := client.Events()
	if err != nil {
		log.Fatal(err)
	}
	for event := range stream.Events {
		if event.Status != "die" {
			continue
		}
		container, err := client.InspectContainer(event.ID)
		if err != nil {
			log.Println("inspect container", event.ID, "error:", err)
			// TODO: set job status anyway?
			continue
		}
		state.SetStatusDone(event.ID, container.State.ExitCode)
	}
}

// TODO: fix this, horribly broken

func allocatePorts(ports chan<- int, startPort, endPort int) {
	for i := startPort; i < endPort; i++ {
		ports <- i
	}
	// TODO: handle wrap-around
}
<commit_msg>host: Merge pull request #12 from flynn/attempt-discoverd-connection<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/flynn/flynn-host/client"
	"github.com/flynn/flynn-host/sampi"
	"github.com/flynn/flynn-host/types"
	"github.com/flynn/go-discoverd"
	"github.com/flynn/go-dockerclient"
	"github.com/flynn/go-flynn/attempt"
	rpc "github.com/flynn/rpcplus/comborpc"
	"github.com/technoweenie/grohl"
)

// Attempts is the attempt strategy that is used to connect to discoverd.
var Attempts = attempt.Strategy{
	Min:   5,
	Total: 5 * time.Second,
	Delay: 200 * time.Millisecond,
}

func main() {
	hostname, _ := os.Hostname()
	externalAddr := flag.String("external", "", "external IP of host")
	configFile := flag.String("config", "", "configuration file")
	manifestFile := flag.String("manifest", "", "manifest file")
	hostID := flag.String("id", hostname, "host id")
	flag.Parse()
	grohl.AddContext("app", "lorne")
	grohl.Log(grohl.Data{"at": "start"})
	g := grohl.NewContext(grohl.Data{"fn": "main"})

	dockerc, err := docker.NewClient("unix:///var/run/docker.sock")
	if err != nil {
		log.Fatal(err)
	}

	state := NewState()
	ports := make(chan int)

	go allocatePorts(ports, 55000, 65535)
	go serveHTTP(&Host{state: state, docker: dockerc}, &attachHandler{state: state, docker: dockerc})
	go streamEvents(dockerc, state)

	processor := &jobProcessor{
		externalAddr: *externalAddr,
		docker:       dockerc,
		state:        state,
		discoverd:    os.Getenv("DISCOVERD"),
	}

	runner := &manifestRunner{
		env:        parseEnviron(),
		externalIP: *externalAddr,
		ports:      ports,
		processor:  processor,
		docker:     dockerc,
	}

	var disc *discoverd.Client
	if *manifestFile != "" {
		f, err := os.Open(*manifestFile)
		if err != nil {
			log.Fatal(err)
		}
		services, err := runner.runManifest(f)
		if err != nil {
			log.Fatal(err)
		}
		f.Close()

		if d, ok := services["discoverd"]; ok {
			processor.discoverd = fmt.Sprintf("%s:%d", d.InternalIP, d.TCPPorts[0])
			var disc *discoverd.Client
			err = Attempts.Run(func() (err error) {
				disc, err = discoverd.NewClientUsingAddress(processor.discoverd)
				return
			})
			if err != nil {
				log.Fatal(err)
			}
		}
	}

	if processor.discoverd == "" && *externalAddr != "" {
		processor.discoverd = *externalAddr + ":1111"
	}
	// HACK: use env as global for discoverd connection in sampic
	os.Setenv("DISCOVERD", processor.discoverd)
	if disc == nil {
		disc, err = discoverd.NewClientUsingAddress(processor.discoverd)
		if err != nil {
			log.Fatal(err)
		}
	}
	sampiStandby, err := disc.RegisterAndStandby("flynn-host", *externalAddr+":1113", map[string]string{"id": *hostID})
	if err != nil {
		log.Fatal(err)
	}
	go func() {
		<-sampiStandby
		rpc.Register(sampi.NewCluster(sampi.NewState()))
	}()

	cluster, err := client.New()
	if err != nil {
		log.Fatal(err)
	}
	g.Log(grohl.Data{"at": "sampi_connected"})

	events := make(chan host.Event)
	state.AddListener("all", events)
	go syncScheduler(cluster, events)

	var h *host.Host
	if *configFile != "" {
		h, err = openConfig(*configFile)
		if err != nil {
			log.Fatal(err)
		}
	} else {
		h = &host.Host{Resources: make(map[string]host.ResourceValue)}
	}
	if _, ok := h.Resources["memory"]; !ok {
		h.Resources["memory"] = host.ResourceValue{Value: 1024}
	}
	h.ID = *hostID
	h.Jobs = state.ClusterJobs()

	jobs := make(chan *host.Job)
	hostErr := cluster.ConnectHost(h, jobs)
	g.Log(grohl.Data{"at": "host_registered"})
	processor.Process(ports, jobs)
	log.Fatal(*hostErr)
}

type jobProcessor struct {
	externalAddr string
	discoverd    string
	docker       interface {
		CreateContainer(*docker.Config) (*docker.Container, error)
		PullImage(docker.PullImageOptions, io.Writer) error
		StartContainer(string, *docker.HostConfig) error
	}
	state *State
}

func (p *jobProcessor) Process(ports <-chan int, jobs chan *host.Job) {
	for job := range jobs {
		p.processJob(ports, job)
	}
}

func (p *jobProcessor) processJob(ports <-chan int, job *host.Job) (*docker.Container, error) {
	g := grohl.NewContext(grohl.Data{"fn": "process_job", "job.id": job.ID})
	g.Log(grohl.Data{"at": "start", "job.image": job.Config.Image, "job.cmd": job.Config.Cmd, "job.entrypoint": job.Config.Entrypoint})

	for i := 0; i < job.TCPPorts; i++ {
		port := strconv.Itoa(<-ports)
		if i == 0 {
			job.Config.Env = append(job.Config.Env, "PORT="+port)
		}
		job.Config.Env = append(job.Config.Env, fmt.Sprintf("PORT_%d=%s", i, port))
		if job.Config.ExposedPorts == nil {
			job.Config.ExposedPorts = make(map[string]struct{}, job.TCPPorts)
		}
		job.Config.ExposedPorts[port+"/tcp"] = struct{}{}
		if job.HostConfig == nil {
			job.HostConfig = &docker.HostConfig{
				PortBindings:    make(map[string][]docker.PortBinding, job.TCPPorts),
				PublishAllPorts: true,
			}
		}
		job.HostConfig.PortBindings[port+"/tcp"] = []docker.PortBinding{{HostPort: port}}
	}
	if p.externalAddr != "" {
		job.Config.Env = appendUnique(job.Config.Env, "EXTERNAL_IP="+p.externalAddr, "SD_HOST="+p.externalAddr, "DISCOVERD="+p.discoverd)
	}
	p.state.AddJob(job)
	g.Log(grohl.Data{"at": "create_container"})
	container, err := p.docker.CreateContainer(job.Config)
	if err == docker.ErrNoSuchImage {
		g.Log(grohl.Data{"at": "pull_image"})
		err = p.docker.PullImage(docker.PullImageOptions{Repository: job.Config.Image}, os.Stdout)
		if err != nil {
			g.Log(grohl.Data{"at": "pull_image", "status": "error", "err": err})
			p.state.SetStatusFailed(job.ID, err)
			return nil, err
		}
		container, err = p.docker.CreateContainer(job.Config)
	}
	if err != nil {
		g.Log(grohl.Data{"at": "create_container", "status": "error", "err": err})
		p.state.SetStatusFailed(job.ID, err)
		return nil, err
	}
	p.state.SetContainerID(job.ID, container.ID)
	p.state.WaitAttach(job.ID)
	g.Log(grohl.Data{"at": "start_container"})
	if err := p.docker.StartContainer(container.ID, job.HostConfig); err != nil {
		g.Log(grohl.Data{"at": "start_container", "status": "error", "err": err})
		p.state.SetStatusFailed(job.ID, err)
		return nil, err
	}
	p.state.SetStatusRunning(job.ID)
	g.Log(grohl.Data{"at": "finish"})
	return container, nil
}

func appendUnique(s []string, vars ...string) []string {
outer:
	for _, v := range vars {
		for _, existing := range s {
			if strings.HasPrefix(existing, strings.SplitN(v, "=", 2)[0]+"=") {
				continue outer
			}
		}
		s = append(s, v)
	}
	return s
}

type sampiSyncClient interface {
	RemoveJobs([]string) error
}

func syncScheduler(scheduler sampiSyncClient, events <-chan host.Event) {
	for event := range events {
		if event.Event != "stop" {
			continue
		}
		grohl.Log(grohl.Data{"fn": "scheduler_event", "at": "remove_job", "job.id": event.JobID})
		if err := scheduler.RemoveJobs([]string{event.JobID}); err != nil {
			grohl.Log(grohl.Data{"fn": "scheduler_event", "at": "remove_job", "status": "error", "err": err, "job.id": event.JobID})
			// TODO: try to reconnect?
		}
	}
}

type dockerStreamClient interface {
	Events() (*docker.EventStream, error)
	InspectContainer(string) (*docker.Container, error)
}

func streamEvents(client dockerStreamClient, state *State) {
	stream, err := client.Events()
	if err != nil {
		log.Fatal(err)
	}
	for event := range stream.Events {
		if event.Status != "die" {
			continue
		}
		container, err := client.InspectContainer(event.ID)
		if err != nil {
			log.Println("inspect container", event.ID, "error:", err)
			// TODO: set job status anyway?
			continue
		}
		state.SetStatusDone(event.ID, container.State.ExitCode)
	}
}

// TODO: fix this, horribly broken

func allocatePorts(ports chan<- int, startPort, endPort int) {
	for i := startPort; i < endPort; i++ {
		ports <- i
	}
	// TODO: handle wrap-around
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"github.com/aschepis/kernctl"
)

const (
	// Providers
	NSTAT_PROVIDER_ROUTE = 1
	NSTAT_PROVIDER_TCP   = 2
	NSTAT_PROVIDER_UDP   = 3

	// generic response messages
	NSTAT_MSG_TYPE_SUCCESS = 0
	NSTAT_MSG_TYPE_ERROR   = 1

	// Requests
	NSTAT_MSG_TYPE_ADD_SRC      = 1001
	NSTAT_MSG_TYPE_ADD_ALL_SRCS = 1002
	NSTAT_MSG_TYPE_REM_SRC      = 1003
	NSTAT_MSG_TYPE_QUERY_SRC    = 1004
	NSTAT_MSG_TYPE_GET_SRC_DESC = 1005

	// Responses/Notfications
	NSTAT_MSG_TYPE_SRC_ADDED   = 10001
	NSTAT_MSG_TYPE_SRC_REMOVED = 10002
	NSTAT_MSG_TYPE_SRC_DESC    = 10003
	NSTAT_MSG_TYPE_SRC_COUNTS  = 10004

	NSTAT_SRC_REF_ALL     = 0xFFFFFFFF
	NSTAT_SRC_REF_INVALID = 0
)

type nstat_msg_hdr struct {
	Context uint64
	HType   uint32
	Pad     uint32 // unused for now
}

/*****************************************************************************/
/* REQUESTS                                                                  */
/*****************************************************************************/

// Type nstat_msg_add_all_srcs, implements kernctl.Message for serialization.
type nstat_msg_add_all_srcs struct {
	Hdr      nstat_msg_hdr
	Provider uint32
}

func (msg *nstat_msg_add_all_srcs) Bytes() []byte {
	buf := new(bytes.Buffer)
	binary.Write(buf, binary.LittleEndian, msg)
	return buf.Bytes()
}

type nstat_msg_query_src_req struct {
	Hdr    nstat_msg_hdr
	SrcRef uint32
}

func (msg *nstat_msg_query_src_req) Bytes() []byte {
	buf := new(bytes.Buffer)
	binary.Write(buf, binary.LittleEndian, msg)
	return buf.Bytes()
}

type nstat_msg_get_src_description struct {
	Hdr    nstat_msg_hdr
	SrcRef uint32
}

func addAllSources(conn *kernctl.Conn, provider uint32) {
	aasreq := &nstat_msg_add_all_srcs{
		Hdr: nstat_msg_hdr{
			HType:   NSTAT_MSG_TYPE_ADD_ALL_SRCS,
			Context: 3,
		},
		Provider: provider,
	}
	fmt.Println("addAllSources", provider)
	conn.SendCommand(aasreq)
}

func getCounts(conn *kernctl.Conn, srcRef uint32) {
	qsreq := &nstat_msg_query_src_req{
		Hdr: nstat_msg_hdr{
			HType:   NSTAT_MSG_TYPE_QUERY_SRC,
			Context: 1005,
		},
		SrcRef: srcRef,
	}
	fmt.Println("getCounts", srcRef)
	conn.SendCommand(qsreq)
}

/*****************************************************************************/
/* RESPONSES                                                                 */
/*****************************************************************************/
type nstat_msg_src_added struct {
	Hdr      nstat_msg_hdr
	Provider uint32
	SrcRef   uint32
}

type nstat_msg_src_removed struct {
	Hdr    nstat_msg_hdr
	SrcRef uint32
}

type nstat_msg_src_description struct {
	Hdr      nstat_msg_hdr
	SrcRef   uint32
	Provider uint32
	// u_int8_t  data[];
}

/*****************************************************************************/

// The original C structures are #pragma pack(1), but here this isn't important,
// we let binary.Read map packed byte stream to unpacked go struct.
type nstat_counts struct {
	Rxpackets         uint64
	Rxbytes           uint64
	Txpackets         uint64
	Txbytes           uint64
	Rxduplicatebytes  uint32
	Rxoutoforderbytes uint32
	Txretransmit      uint32
	Connectattempts   uint32
	Connectsuccesses  uint32
	Min_rtt           uint32
	Avg_rtt           uint32
	Var_rtt           uint32
}

type nstat_msg_src_counts struct {
	Hdr    nstat_msg_hdr
	SrcRef uint32
	Counts nstat_counts
}

/* Process the response we received from the system socket. */
func process_nstat_msg(msg_hdr nstat_msg_hdr, buf []byte) error {

	switch msg_hdr.HType {
	case NSTAT_MSG_TYPE_SRC_ADDED:
		var msg nstat_msg_src_added
		reader := bytes.NewReader(buf)
		err := binary.Read(reader, binary.LittleEndian, &msg)
		if err != nil {
			fmt.Println("binary.Read failed:", err)
			break
		}
		fmt.Println("new source: ", msg)

	case NSTAT_MSG_TYPE_SRC_REMOVED:
		var msg nstat_msg_src_removed
		reader := bytes.NewReader(buf)
		err := binary.Read(reader, binary.LittleEndian, &msg)
		if err != nil {
			fmt.Println("binary.Read failed:", err)
			break
		}
		fmt.Println("source removed: ", msg)

	case NSTAT_MSG_TYPE_SRC_DESC:
		var msg nstat_msg_src_description
		reader := bytes.NewReader(buf)
		err := binary.Read(reader, binary.LittleEndian, &msg)
		if err != nil {
			fmt.Println("binary.Read src_description failed:", err)
			break
		}
		switch msg.Provider {
		case NSTAT_PROVIDER_TCP:
			fmt.Println("TCP description received: ", msg)
		case NSTAT_PROVIDER_UDP:
			fmt.Println("UDP description received: ", msg)
		}
		fmt.Println("description received: ", msg)

	case NSTAT_MSG_TYPE_SRC_COUNTS:
		var msg nstat_msg_src_counts
		reader := bytes.NewReader(buf)
		err := binary.Read(reader, binary.LittleEndian, &msg)
		if err != nil {
			fmt.Println("binary.Read failed:", err)
			break
		}
		fmt.Println("counts received: ", msg)

	}
	return nil
}

const (
	STATE_INITIAL      = 0
	STATE_TCP_ADDED    = 2
	STATE_UDP_ADDED    = 4
	STATE_COUNTS_ADDED = 6
)

func main() {
	conn := kernctl.NewConnByName("com.apple.network.statistics")
	if err := conn.Connect(); err != nil {
		panic(err)
	}

	var state = STATE_INITIAL
	for {
		// Subscribe to following events one by one:
		// 1. all TCP events
		// 2. all UDP events
		// 3. XXXXXXX
		//
		switch state {
		case STATE_INITIAL:
			addAllSources(conn, NSTAT_PROVIDER_TCP)
			state++
		case STATE_TCP_ADDED:
			addAllSources(conn, NSTAT_PROVIDER_UDP)
			state++
		case STATE_UDP_ADDED:
			getCounts(conn, NSTAT_SRC_REF_ALL)
			state++
		default:
			/* in one of the waiting states (uneven numbers) */
			break
		}

		if err, buf := conn.Select(2048); err != nil {
			panic(err)
		} else {
			var msg_hdr nstat_msg_hdr

			reader := bytes.NewReader(buf)
			err := binary.Read(reader, binary.LittleEndian, &msg_hdr)
			if err != nil {
				fmt.Println("binary.Read failed:", err)
				break
			}
			fmt.Println("msg_hdr recvd:", msg_hdr)

			switch msg_hdr.HType {
			case NSTAT_MSG_TYPE_SUCCESS:
				{
					/* Previous requested action was successful, go to next. */
					state++
				}
			case NSTAT_MSG_TYPE_SRC_ADDED, NSTAT_MSG_TYPE_SRC_REMOVED,
				NSTAT_MSG_TYPE_SRC_DESC, NSTAT_MSG_TYPE_SRC_COUNTS:
				{
					ret := process_nstat_msg(msg_hdr, buf)
					if ret != nil {
						break
					}
				}
			}

		}
	}

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

package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"github.com/aschepis/kernctl"
)

const (
	// Providers
	NSTAT_PROVIDER_ROUTE = 1
	NSTAT_PROVIDER_TCP   = 2
	NSTAT_PROVIDER_UDP   = 3

	// generic response messages
	NSTAT_MSG_TYPE_SUCCESS = 0
	NSTAT_MSG_TYPE_ERROR   = 1

	// Requests
	NSTAT_MSG_TYPE_ADD_SRC      = 1001
	NSTAT_MSG_TYPE_ADD_ALL_SRCS = 1002
	NSTAT_MSG_TYPE_REM_SRC      = 1003
	NSTAT_MSG_TYPE_QUERY_SRC    = 1004
	NSTAT_MSG_TYPE_GET_SRC_DESC = 1005

	// Responses/Notfications
	NSTAT_MSG_TYPE_SRC_ADDED   = 10001
	NSTAT_MSG_TYPE_SRC_REMOVED = 10002
	NSTAT_MSG_TYPE_SRC_DESC    = 10003
	NSTAT_MSG_TYPE_SRC_COUNTS  = 10004

	NSTAT_SRC_REF_ALL     = 0xFFFFFFFF
	NSTAT_SRC_REF_INVALID = 0
)

type nstat_msg_hdr struct {
	Context uint64
	HType   uint32
	Pad     uint32 // unused for now
}

/*****************************************************************************/
/* REQUESTS                                                                  */
/*****************************************************************************/

// Type nstat_msg_add_all_srcs, implements kernctl.Message for serialization.
type nstat_msg_add_all_srcs struct {
	Hdr      nstat_msg_hdr
	Provider uint32
}

func (msg *nstat_msg_add_all_srcs) Bytes() []byte {
	buf := new(bytes.Buffer)
	binary.Write(buf, binary.LittleEndian, msg)
	return buf.Bytes()
}

type nstat_msg_query_src_req struct {
	Hdr    nstat_msg_hdr
	SrcRef uint32
}

func (msg *nstat_msg_query_src_req) Bytes() []byte {
	buf := new(bytes.Buffer)
	binary.Write(buf, binary.LittleEndian, msg)
	return buf.Bytes()
}

type nstat_msg_get_src_description struct {
	Hdr    nstat_msg_hdr
	SrcRef uint32
}

func addAllSources(conn *kernctl.Conn, provider uint32) {
	aasreq := &nstat_msg_add_all_srcs{
		Hdr: nstat_msg_hdr{
			HType:   NSTAT_MSG_TYPE_ADD_ALL_SRCS,
			Context: 3,
		},
		Provider: provider,
	}
	fmt.Println("addAllSources", provider)
	conn.SendCommand(aasreq)
}

func getCounts(conn *kernctl.Conn, srcRef uint32) {
	qsreq := &nstat_msg_query_src_req{
		Hdr: nstat_msg_hdr{
			HType:   NSTAT_MSG_TYPE_QUERY_SRC,
			Context: 1005,
		},
		SrcRef: srcRef,
	}
	fmt.Println("getCounts", srcRef)
	conn.SendCommand(qsreq)
}

/*****************************************************************************/
/* RESPONSES                                                                 */
/*****************************************************************************/
type nstat_msg_src_added struct {
	Hdr      nstat_msg_hdr
	Provider uint32
	SrcRef   uint32
}

type nstat_msg_src_removed struct {
	Hdr    nstat_msg_hdr
	SrcRef uint32
}

type nstat_msg_src_description struct {
	Hdr      nstat_msg_hdr
	SrcRef   uint32
	Provider uint32
	// u_int8_t  data[];
}

/*****************************************************************************/

// The original C structures are #pragma pack(1), but here this isn't important,
// we let binary.Read map packed byte stream to unpacked go struct.
type nstat_counts struct {
	Rxpackets         uint64
	Rxbytes           uint64
	Txpackets         uint64
	Txbytes           uint64
	Rxduplicatebytes  uint32
	Rxoutoforderbytes uint32
	Txretransmit      uint32
	Connectattempts   uint32
	Connectsuccesses  uint32
	Min_rtt           uint32
	Avg_rtt           uint32
	Var_rtt           uint32
}

type nstat_msg_src_counts struct {
	Hdr    nstat_msg_hdr
	SrcRef uint32
	Counts nstat_counts
}

/* Process the response we received from the system socket. */
func process_nstat_msg(msg_hdr nstat_msg_hdr, buf []byte) error {

	switch msg_hdr.HType {
	case NSTAT_MSG_TYPE_SRC_ADDED:
		var msg nstat_msg_src_added
		reader := bytes.NewReader(buf)
		err := binary.Read(reader, binary.LittleEndian, &msg)
		if err != nil {
			fmt.Println("binary.Read failed:", err)
			break
		}
		fmt.Println("new source: ", msg)

	case NSTAT_MSG_TYPE_SRC_REMOVED:
		var msg nstat_msg_src_removed
		reader := bytes.NewReader(buf)
		err := binary.Read(reader, binary.LittleEndian, &msg)
		if err != nil {
			fmt.Println("binary.Read failed:", err)
			break
		}
		fmt.Println("source removed: ", msg)

	case NSTAT_MSG_TYPE_SRC_DESC:
		var msg nstat_msg_src_description
		reader := bytes.NewReader(buf)
		err := binary.Read(reader, binary.LittleEndian, &msg)
		if err != nil {
			fmt.Println("binary.Read src_description failed:", err)
			break
		}
		switch msg.Provider {
		case NSTAT_PROVIDER_TCP:
			fmt.Println("TCP description received: ", msg)
		case NSTAT_PROVIDER_UDP:
			fmt.Println("UDP description received: ", msg)
		}
		fmt.Println("description received: ", msg)

	case NSTAT_MSG_TYPE_SRC_COUNTS:
		var msg nstat_msg_src_counts
		reader := bytes.NewReader(buf)
		err := binary.Read(reader, binary.LittleEndian, &msg)
		if err != nil {
			fmt.Println("binary.Read failed:", err)
			break
		}
		fmt.Println("counts received: ", msg)

	}
	return nil
}

const (
	STATE_INITIAL      = 0
	STATE_TCP_ADDED    = 2
	STATE_UDP_ADDED    = 4
	STATE_COUNTS_ADDED = 6
)

func main() {
	conn := kernctl.NewConnByName("com.apple.network.statistics")
	if err := conn.Connect(); err != nil {
		panic(err)
	}

	var state = STATE_INITIAL
	for {
		// Subscribe to following events one by one:
		// 1. all TCP events
		// 2. all UDP events
		// 3. XXXXXXX
		//
		switch state {
		case STATE_INITIAL:
			addAllSources(conn, NSTAT_PROVIDER_TCP)
			state++
		case STATE_TCP_ADDED:
			addAllSources(conn, NSTAT_PROVIDER_UDP)
			state++
		case STATE_UDP_ADDED:
			getCounts(conn, NSTAT_SRC_REF_ALL)
			state++
		default:
			/* in one of the waiting states (uneven numbers) */
			break
		}

		if err, buf := conn.Select(2048); err != nil {
			panic(err)
		} else {
			var msg_hdr nstat_msg_hdr

			reader := bytes.NewReader(buf)
			err := binary.Read(reader, binary.LittleEndian, &msg_hdr)
			if err != nil {
				fmt.Println("binary.Read failed:", err)
				break
			}
			fmt.Println("msg_hdr recvd:", msg_hdr)

			switch msg_hdr.HType {
			case NSTAT_MSG_TYPE_SUCCESS:
				{
					/* Previous requested action was successful, go to next. */
					state++
				}
			case NSTAT_MSG_TYPE_SRC_ADDED, NSTAT_MSG_TYPE_SRC_REMOVED,
				NSTAT_MSG_TYPE_SRC_DESC, NSTAT_MSG_TYPE_SRC_COUNTS:
				{
					ret := process_nstat_msg(msg_hdr, buf)
					if ret != nil {
						break
					}
				}
			}

		}
	}

	conn.Close()
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 CHEN Xianren. All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package tiny

import (
	"net/http"
	"reflect"
	"sync"

	"github.com/cxr29/log"
)

type Context struct {
	http.ResponseWriter
	Request       *http.Request
	Values        map[interface{}]interface{}
	Params        []string
	tree          *Tree
	node          *Node
	wroteHeader   bool
	written       int64
	status, index int
}

func (ctx *Context) Routed() bool {
	return ctx.node != nil
}

func (ctx *Context) call(i int) {
	var h Handler
	if j := i - len(ctx.tree.handlers); j < 0 {
		h = ctx.tree.handlers[i]
	} else if ctx.Routed() && j < len(ctx.node.handlers) {
		h = ctx.node.handlers[j]
	} else {
		return
	}
	h.ServeHTTP(ctx)
	if ctx.index == i && !ctx.wroteHeader {
		ctx.Next()
	}
}

func (ctx *Context) Next() {
	ctx.index++
	ctx.call(ctx.index)
}

func (ctx *Context) WriteHeader(code int) {
	if !ctx.wroteHeader {
		ctx.wroteHeader = true
		ctx.status = code
	}
	ctx.ResponseWriter.WriteHeader(code)
}

const contentType = "Content-Type"

func (ctx *Context) Write(data []byte) (int, error) {
	if !ctx.wroteHeader {
		h := ctx.Header()
		if h.Get("Transfer-Encoding") == "" && h.Get(contentType) == "" {
			h.Set(contentType, http.DetectContentType(data))
		}
		ctx.WriteHeader(http.StatusOK)
	}
	n, err := ctx.ResponseWriter.Write(data)
	log.ErrWarning(err)
	ctx.written += int64(n)
	return n, err
}

func (ctx *Context) WroteHeader() bool {
	return ctx.wroteHeader
}

func (ctx *Context) Written() int64 {
	return ctx.written
}

func (ctx *Context) Status() int {
	return ctx.status
}

func (ctx *Context) Param(name string) string {
	if ctx.Routed() {
		for i := len(ctx.node.params) - 1; i >= 0; i-- {
			if ctx.node.params[i] == name {
				return ctx.Params[i]
			}
		}
	}
	return ""
}

type valueKey struct {
	id uint
}

var (
	id uint
	mu sync.Mutex
)

func NewValueKey() (key valueKey) {
	mu.Lock()
	id++
	key.id = id
	mu.Unlock()
	return
}

func (ctx *Context) Value(k interface{}) interface{} {
	if v, ok := ctx.Values[k]; ok {
		return v
	} else {
		panic("key not exist")
	}
}

func (ctx *Context) SetValue(k, v interface{}) {
	if k == nil {
		panic("nil key")
	} else if !reflect.TypeOf(k).Comparable() {
		panic("not comparable key")
	} else if ctx.Values == nil {
		ctx.Values = map[interface{}]interface{}{k: v}
	} else if _, ok := ctx.Values[k]; ok {
		panic("key already exists")
	} else {
		ctx.Values[k] = v
	}
}
<commit_msg>make ten values<commit_after>// Copyright (c) 2016 CHEN Xianren. All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package tiny

import (
	"net/http"
	"reflect"
	"sync"

	"github.com/cxr29/log"
)

type Context struct {
	http.ResponseWriter
	Request       *http.Request
	Values        map[interface{}]interface{}
	Params        []string
	tree          *Tree
	node          *Node
	wroteHeader   bool
	written       int64
	status, index int
}

func (ctx *Context) Routed() bool {
	return ctx.node != nil
}

func (ctx *Context) call(i int) {
	var h Handler
	if j := i - len(ctx.tree.handlers); j < 0 {
		h = ctx.tree.handlers[i]
	} else if ctx.Routed() && j < len(ctx.node.handlers) {
		h = ctx.node.handlers[j]
	} else {
		return
	}
	h.ServeHTTP(ctx)
	if ctx.index == i && !ctx.wroteHeader {
		ctx.Next()
	}
}

func (ctx *Context) Next() {
	ctx.index++
	ctx.call(ctx.index)
}

func (ctx *Context) WriteHeader(code int) {
	if !ctx.wroteHeader {
		ctx.wroteHeader = true
		ctx.status = code
	}
	ctx.ResponseWriter.WriteHeader(code)
}

const contentType = "Content-Type"

func (ctx *Context) Write(data []byte) (int, error) {
	if !ctx.wroteHeader {
		h := ctx.Header()
		if h.Get("Transfer-Encoding") == "" && h.Get(contentType) == "" {
			h.Set(contentType, http.DetectContentType(data))
		}
		ctx.WriteHeader(http.StatusOK)
	}
	n, err := ctx.ResponseWriter.Write(data)
	log.ErrWarning(err)
	ctx.written += int64(n)
	return n, err
}

func (ctx *Context) WroteHeader() bool {
	return ctx.wroteHeader
}

func (ctx *Context) Written() int64 {
	return ctx.written
}

func (ctx *Context) Status() int {
	return ctx.status
}

func (ctx *Context) Param(name string) string {
	if ctx.Routed() {
		for i := len(ctx.node.params) - 1; i >= 0; i-- {
			if ctx.node.params[i] == name {
				return ctx.Params[i]
			}
		}
	}
	return ""
}

type valueKey struct {
	id uint
}

var (
	id uint
	mu sync.Mutex
)

func NewValueKey() (key valueKey) {
	mu.Lock()
	id++
	key.id = id
	mu.Unlock()
	return
}

func (ctx *Context) Value(k interface{}) interface{} {
	if v, ok := ctx.Values[k]; ok {
		return v
	} else {
		panic("key not exist")
	}
}

func (ctx *Context) SetValue(k, v interface{}) {
	if k == nil {
		panic("nil key")
	} else if !reflect.TypeOf(k).Comparable() {
		panic("not comparable key")
	} else if ctx.Values == nil {
		ctx.Values = make(map[interface{}]interface{}, 10)
	} else if _, ok := ctx.Values[k]; ok {
		panic("key already exists")
	}
	ctx.Values[k] = v
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"github.com/powerman/rpc-codec/jsonrpc2"
	"golang.org/x/net/websocket"
	"net/rpc"
	"strings"
	"time"
	"unicode"
)

type Context struct {
	// There is a timeout after releasing a stream during which it is possible
	// to reacquire the same object and continue broadcasting. Once the timeout
	// elapses, the stream is closed for good.
	Timeout time.Duration
	// How many messages to transmit on a `Chat.RequestHistory` RPC call.
	ChatHistory int

	streams map[string]*BroadcastContext
}

type BroadcastContext struct {
	Broadcast
	closing            bool
	closingStateChange chan bool
	chatters           map[*chatterContext]int // A hash set. Values are ignored.
	chattersNames      map[string]int          // Same.
	chatHistory        chatMessageQueue
	// These values are for the whole stream, so they include audio and muxing overhead.
	// The latter is negligible, however, and the former is normally about 64k,
	// so also negligible. Or at least predictable.
	RateMean float64
	RateVar  float64
	rateUnit float64
}

type chatterContext struct {
	name   string
	socket *websocket.Conn
	stream *BroadcastContext
}

type chatMessage struct {
	name string
	text string
}

type chatMessageQueue struct {
	data  []chatMessage
	start int
}

func newChatMessageQueue(size int) chatMessageQueue {
	return chatMessageQueue{make([]chatMessage, 0, size), 0}
}

func (q *chatMessageQueue) Push(x chatMessage) {
	if len(q.data) == cap(q.data) {
		q.data[q.start] = x
		q.start = (q.start + 1) % len(q.data)
	} else {
		q.data = q.data[:len(q.data)+1]
		q.data[len(q.data)-1] = x
	}
}

func (q *chatMessageQueue) Iterate(f func(x chatMessage) error) error {
	for i := 0; i < len(q.data); i++ {
		if err := f(q.data[(i+q.start)%len(q.data)]); err != nil {
			return err
		}
	}
	return nil
}

// Acquire a stream for writing. Only one "writable" reference can be held;
// until it is closed, this function will return an error.
func (ctx *Context) Acquire(id string) (*BroadcastContext, bool) {
	if ctx.streams == nil {
		ctx.streams = make(map[string]*BroadcastContext)
	}
	stream, ok := ctx.streams[id]
	if !ok {
		v := BroadcastContext{
			Broadcast:          NewBroadcast(),
			closingStateChange: make(chan bool),
			chatters:           make(map[*chatterContext]int),
			chattersNames:      make(map[string]int),
			chatHistory:        newChatMessageQueue(ctx.ChatHistory),
		}
		ctx.streams[id] = &v
		go v.monitor(ctx, id)
		return &v, true
	}
	if !stream.closing {
		return nil, false
	}
	stream.closingStateChange <- false
	return stream, true
}

// Acquire a stream for reading. There is no limit on the number of concurrent readers.
func (ctx *Context) Get(id string) (*BroadcastContext, bool) {
	if ctx.streams == nil {
		return nil, false
	}
	stream, ok := ctx.streams[id]
	return stream, ok
}

func (stream *BroadcastContext) monitor(ctx *Context, id string) {
	ticker := time.NewTicker(time.Second)
	ticksWhileOffline := 0 * time.Second

	for {
		select {
		case stream.closing = <-stream.closingStateChange:
			ticksWhileOffline = 0
		case <-ticker.C:
			if stream.closing {
				if ticksWhileOffline += time.Second; ticksWhileOffline > ctx.Timeout {
					delete(ctx.streams, id)
					ticker.Stop()
					stream.Broadcast.Close()
					for ctx := range stream.chatters {
						ctx.socket.Close()
					}
					return
				}
			}
			// exponentially weighted moving moments at a = 0.5
			//     avg[n] = a * x + (1 - a) * avg[n - 1]
			//     var[n] = a * (x - avg[n]) ** 2 / (1 - a) + (1 - a) * var[n - 1]
			stream.RateMean += stream.rateUnit / 2
			stream.RateVar += stream.rateUnit*stream.rateUnit - stream.RateVar/2
			stream.rateUnit = -stream.RateMean
		}
	}
}

func (stream *BroadcastContext) Write(data []byte) (int, error) {
	stream.rateUnit += float64(len(data))
	sent, err := stream.Broadcast.Write(data)
	return sent, err
}

func (stream *BroadcastContext) Close() {
	stream.closingStateChange <- true
}

type RPCSingleStringArg struct {
	First string
}

func (x *RPCSingleStringArg) UnmarshalJSON(buf []byte) error {
	fields := []interface{}{&x.First}
	expect := len(fields)
	if err := json.Unmarshal(buf, &fields); err != nil {
		return err
	}
	if len(fields) != expect {
		return errors.New("invalid number of arguments")
	}
	return nil
}

func (ctx *chatterContext) SetName(args *RPCSingleStringArg, _ *interface{}) error {
	name := strings.TrimSpace(args.First)
	if len(name) == 0 {
		return errors.New("name must not be empty")
	}
	if len(name) > 32 {
		return errors.New("name too long")
	}
	for _, c := range name {
		if !unicode.IsGraphic(c) {
			return errors.New("name contains invalid characters")
		}
	}
	if _, ok := ctx.stream.chattersNames[name]; ok {
		return errors.New("name already taken")
	}
	ctx.stream.chattersNames[name] = 0
	if ctx.name != "" {
		delete(ctx.stream.chattersNames, ctx.name)
	}
	ctx.name = name
	return nil
}

func (ctx *chatterContext) SendMessage(args *RPCSingleStringArg, _ *interface{}) error {
	if ctx.name == "" {
		return errors.New("must obtain a name first")
	}
	msg := chatMessage{ctx.name, strings.TrimSpace(args.First)}
	if len(msg.text) == 0 || len(msg.text) > 256 {
		return errors.New("message must have between 1 and 256 characters")
	}
	for viewer := range ctx.stream.chatters {
		viewer.pushMessage(msg)
	}
	ctx.stream.chatHistory.Push(msg)
	return nil
}

func (ctx *chatterContext) RequestHistory(_ *interface{}, _ *interface{}) error {
	return ctx.stream.chatHistory.Iterate(ctx.pushMessage)
}

func (ctx *chatterContext) pushMessage(msg chatMessage) error {
	return pushEvent(ctx.socket, "Chat.Message", []interface{}{msg.name, msg.text})
}

func pushEvent(ws *websocket.Conn, name string, args []interface{}) error {
	return websocket.JSON.Send(ws, map[string]interface{}{
		"jsonrpc": "2.0", "method": name, "params": args,
	})
}

func (stream *BroadcastContext) RunRPC(ws *websocket.Conn) {
	chatter := chatterContext{name: "", socket: ws, stream: stream}
	stream.chatters[&chatter] = 0
	defer func() {
		delete(stream.chatters, &chatter)
		if chatter.name != "" {
			delete(stream.chattersNames, chatter.name)
		}
	}()

	server := rpc.NewServer()
	server.RegisterName("Chat", &chatter)
	server.ServeCodec(jsonrpc2.NewServerCodec(ws, server))
}
<commit_msg>Track viewer count & record stream creation time.<commit_after>package main

import (
	"encoding/json"
	"errors"
	"github.com/powerman/rpc-codec/jsonrpc2"
	"golang.org/x/net/websocket"
	"net/rpc"
	"strings"
	"time"
	"unicode"
)

type Context struct {
	// There is a timeout after releasing a stream during which it is possible
	// to reacquire the same object and continue broadcasting. Once the timeout
	// elapses, the stream is closed for good.
	Timeout time.Duration
	// How many messages to transmit on a `Chat.RequestHistory` RPC call.
	ChatHistory int

	streams map[string]*BroadcastContext
}

type BroadcastContext struct {
	Broadcast
	Created            time.Time
	Viewers            uint
	closing            bool
	closingStateChange chan bool
	chatters           map[*chatterContext]int // A hash set. Values are ignored.
	chattersNames      map[string]int          // Same.
	chatHistory        chatMessageQueue
	// These values are for the whole stream, so they include audio and muxing overhead.
	// The latter is negligible, however, and the former is normally about 64k,
	// so also negligible. Or at least predictable.
	RateMean float64
	RateVar  float64
	rateUnit float64
}

type chatterContext struct {
	name   string
	socket *websocket.Conn
	stream *BroadcastContext
}

type chatMessage struct {
	name string
	text string
}

type chatMessageQueue struct {
	data  []chatMessage
	start int
}

func newChatMessageQueue(size int) chatMessageQueue {
	return chatMessageQueue{make([]chatMessage, 0, size), 0}
}

func (q *chatMessageQueue) Push(x chatMessage) {
	if len(q.data) == cap(q.data) {
		q.data[q.start] = x
		q.start = (q.start + 1) % len(q.data)
	} else {
		q.data = q.data[:len(q.data)+1]
		q.data[len(q.data)-1] = x
	}
}

func (q *chatMessageQueue) Iterate(f func(x chatMessage) error) error {
	for i := 0; i < len(q.data); i++ {
		if err := f(q.data[(i+q.start)%len(q.data)]); err != nil {
			return err
		}
	}
	return nil
}

// Acquire a stream for writing. Only one "writable" reference can be held;
// until it is closed, this function will return an error.
func (ctx *Context) Acquire(id string) (*BroadcastContext, bool) {
	if ctx.streams == nil {
		ctx.streams = make(map[string]*BroadcastContext)
	}
	stream, ok := ctx.streams[id]
	if !ok {
		v := BroadcastContext{
			Broadcast:          NewBroadcast(),
			Created:            time.Now().UTC(),
			closingStateChange: make(chan bool),
			chatters:           make(map[*chatterContext]int),
			chattersNames:      make(map[string]int),
			chatHistory:        newChatMessageQueue(ctx.ChatHistory),
		}
		ctx.streams[id] = &v
		go v.monitor(ctx, id)
		return &v, true
	}
	if !stream.closing {
		return nil, false
	}
	stream.closingStateChange <- false
	return stream, true
}

// Acquire a stream for reading. There is no limit on the number of concurrent readers.
func (ctx *Context) Get(id string) (*BroadcastContext, bool) {
	if ctx.streams == nil {
		return nil, false
	}
	stream, ok := ctx.streams[id]
	return stream, ok
}

func (stream *BroadcastContext) monitor(ctx *Context, id string) {
	ticker := time.NewTicker(time.Second)
	ticksWhileOffline := 0 * time.Second

	for {
		select {
		case stream.closing = <-stream.closingStateChange:
			ticksWhileOffline = 0
		case <-ticker.C:
			if stream.closing {
				if ticksWhileOffline += time.Second; ticksWhileOffline > ctx.Timeout {
					delete(ctx.streams, id)
					ticker.Stop()
					stream.Broadcast.Close()
					for ctx := range stream.chatters {
						ctx.socket.Close()
					}
					return
				}
			}
			// exponentially weighted moving moments at a = 0.5
			//     avg[n] = a * x + (1 - a) * avg[n - 1]
			//     var[n] = a * (x - avg[n]) ** 2 / (1 - a) + (1 - a) * var[n - 1]
			stream.RateMean += stream.rateUnit / 2
			stream.RateVar += stream.rateUnit*stream.rateUnit - stream.RateVar/2
			stream.rateUnit = -stream.RateMean
		}
	}
}

func (stream *BroadcastContext) Write(data []byte) (int, error) {
	stream.rateUnit += float64(len(data))
	sent, err := stream.Broadcast.Write(data)
	return sent, err
}

func (stream *BroadcastContext) Close() {
	stream.closingStateChange <- true
}

type RPCSingleStringArg struct {
	First string
}

func (x *RPCSingleStringArg) UnmarshalJSON(buf []byte) error {
	fields := []interface{}{&x.First}
	expect := len(fields)
	if err := json.Unmarshal(buf, &fields); err != nil {
		return err
	}
	if len(fields) != expect {
		return errors.New("invalid number of arguments")
	}
	return nil
}

func (ctx *chatterContext) SetName(args *RPCSingleStringArg, _ *interface{}) error {
	name := strings.TrimSpace(args.First)
	if len(name) == 0 {
		return errors.New("name must not be empty")
	}
	if len(name) > 32 {
		return errors.New("name too long")
	}
	for _, c := range name {
		if !unicode.IsGraphic(c) {
			return errors.New("name contains invalid characters")
		}
	}
	if _, ok := ctx.stream.chattersNames[name]; ok {
		return errors.New("name already taken")
	}
	ctx.stream.chattersNames[name] = 0
	if ctx.name != "" {
		delete(ctx.stream.chattersNames, ctx.name)
	}
	ctx.name = name
	return nil
}

func (ctx *chatterContext) SendMessage(args *RPCSingleStringArg, _ *interface{}) error {
	if ctx.name == "" {
		return errors.New("must obtain a name first")
	}
	msg := chatMessage{ctx.name, strings.TrimSpace(args.First)}
	if len(msg.text) == 0 || len(msg.text) > 256 {
		return errors.New("message must have between 1 and 256 characters")
	}
	for viewer := range ctx.stream.chatters {
		viewer.pushMessage(msg)
	}
	ctx.stream.chatHistory.Push(msg)
	return nil
}

func (ctx *chatterContext) RequestHistory(_ *interface{}, _ *interface{}) error {
	return ctx.stream.chatHistory.Iterate(ctx.pushMessage)
}

func (ctx *chatterContext) pushMessage(msg chatMessage) error {
	return pushEvent(ctx.socket, "Chat.Message", []interface{}{msg.name, msg.text})
}

func pushEvent(ws *websocket.Conn, name string, args []interface{}) error {
	return websocket.JSON.Send(ws, map[string]interface{}{
		"jsonrpc": "2.0", "method": name, "params": args,
	})
}

func (stream *BroadcastContext) RunRPC(ws *websocket.Conn) {
	chatter := chatterContext{name: "", socket: ws, stream: stream}
	stream.Viewers += 1
	stream.chatters[&chatter] = 0
	defer func() {
		stream.Viewers -= 1
		delete(stream.chatters, &chatter)
		if chatter.name != "" {
			delete(stream.chattersNames, chatter.name)
		}
	}()

	server := rpc.NewServer()
	server.RegisterName("Chat", &chatter)
	server.ServeCodec(jsonrpc2.NewServerCodec(ws, server))
}
<|endoftext|>


<commit_before>package gin

import (
	"bytes"
	"errors"
	"fmt"
	"github.com/gin-gonic/gin/binding"
	"github.com/gin-gonic/gin/render"
	"github.com/julienschmidt/httprouter"
	"log"
	"net/http"
)

const (
	ErrorTypeInternal = 1 << iota
	ErrorTypeExternal = 1 << iota
	ErrorTypeAll      = 0xffffffff
)

// Used internally to collect errors that occurred during an http request.
type errorMsg struct {
	Err  string      `json:"error"`
	Type uint32      `json:"-"`
	Meta interface{} `json:"meta"`
}

type errorMsgs []errorMsg

func (a errorMsgs) ByType(typ uint32) errorMsgs {
	if len(a) == 0 {
		return a
	}
	result := make(errorMsgs, 0, len(a))
	for _, msg := range a {
		if msg.Type&typ > 0 {
			result = append(result, msg)
		}
	}
	return result
}

func (a errorMsgs) String() string {
	if len(a) == 0 {
		return ""
	}
	var buffer bytes.Buffer
	for i, msg := range a {
		text := fmt.Sprintf("Error #%02d: %s \n     Meta: %v\n", (i + 1), msg.Err, msg.Meta)
		buffer.WriteString(text)
	}
	return buffer.String()
}

// Context is the most important part of gin. It allows us to pass variables between middleware,
// manage the flow, validate the JSON of a request and render a JSON response for example.
type Context struct {
	writermem responseWriter
	Request   *http.Request
	Writer    ResponseWriter
	Keys      map[string]interface{}
	Errors    errorMsgs
	Params    httprouter.Params
	Engine    *Engine
	handlers  []HandlerFunc
	index     int8
}

/************************************/
/********** ROUTES GROUPING *********/
/************************************/

func (engine *Engine) createContext(w http.ResponseWriter, req *http.Request, params httprouter.Params, handlers []HandlerFunc) *Context {
	c := engine.cache.Get().(*Context)
	c.writermem.reset(w)
	c.Request = req
	c.Params = params
	c.handlers = handlers
	c.Keys = nil
	c.index = -1
	c.Errors = c.Errors[0:0]
	return c
}

/************************************/
/****** FLOW AND ERROR MANAGEMENT****/
/************************************/

func (c *Context) Copy() *Context {
	var cp Context = *c
	cp.index = AbortIndex
	cp.handlers = nil
	return &cp
}

// Next should be used only in the middlewares.
// It executes the pending handlers in the chain inside the calling handler.
// See example in github.
func (c *Context) Next() {
	c.index++
	s := int8(len(c.handlers))
	for ; c.index < s; c.index++ {
		c.handlers[c.index](c)
	}
}

// Forces the system to do not continue calling the pending handlers.
// For example, the first handler checks if the request is authorized. If it's not, context.Abort(401) should be called.
// The rest of pending handlers would never be called for that request.
func (c *Context) Abort(code int) {
	if code >= 0 {
		c.Writer.WriteHeader(code)
	}
	c.index = AbortIndex
}

// Fail is the same as Abort plus an error message.
// Calling `context.Fail(500, err)` is equivalent to:
// ```
// context.Error("Operation aborted", err)
// context.Abort(500)
// ```
func (c *Context) Fail(code int, err error) {
	c.Error(err, "Operation aborted")
	c.Abort(code)
}

func (c *Context) ErrorTyped(err error, typ uint32, meta interface{}) {
	c.Errors = append(c.Errors, errorMsg{
		Err:  err.Error(),
		Type: typ,
		Meta: meta,
	})
}

// Attaches an error to the current context. The error is pushed to a list of errors.
// It's a good idea to call Error for each error that occurred during the resolution of a request.
// A middleware can be used to collect all the errors and push them to a database together, print a log, or append it in the HTTP response.
func (c *Context) Error(err error, meta interface{}) {
	c.ErrorTyped(err, ErrorTypeExternal, meta)
}

func (c *Context) LastError() error {
	s := len(c.Errors)
	if s > 0 {
		return errors.New(c.Errors[s-1].Err)
	} else {
		return nil
	}
}

/************************************/
/******** METADATA MANAGEMENT********/
/************************************/

// Sets a new pair key/value just for the specified context.
// It also lazy initializes the hashmap.
func (c *Context) Set(key string, item interface{}) {
	if c.Keys == nil {
		c.Keys = make(map[string]interface{})
	}
	c.Keys[key] = item
}

// Get returns the value for the given key or an error if the key does not exist.
func (c *Context) Get(key string) (interface{}, error) {
	if c.Keys != nil {
		item, ok := c.Keys[key]
		if ok {
			return item, nil
		}
	}
	return nil, errors.New("Key does not exist.")
}

// MustGet returns the value for the given key or panics if the value doesn't exist.
func (c *Context) MustGet(key string) interface{} {
	value, err := c.Get(key)
	if err != nil || value == nil {
		log.Panicf("Key %s doesn't exist", key)
	}
	return value
}

/************************************/
/******** ENCOGING MANAGEMENT********/
/************************************/

// This function checks the Content-Type to select a binding engine automatically,
// Depending the "Content-Type" header different bindings are used:
// "application/json" --> JSON binding
// "application/xml"  --> XML binding
// else --> returns an error
// if Parses the request's body as JSON if Content-Type == "application/json"  using JSON or XML  as a JSON input. It decodes the json payload into the struct specified as a pointer.Like ParseBody() but this method also writes a 400 error if the json is not valid.
func (c *Context) Bind(obj interface{}) bool {
	var b binding.Binding
	ctype := filterFlags(c.Request.Header.Get("Content-Type"))
	switch {
	case c.Request.Method == "GET" || ctype == MIMEPOSTForm:
		b = binding.Form
	case ctype == MIMEJSON:
		b = binding.JSON
	case ctype == MIMEXML || ctype == MIMEXML2:
		b = binding.XML
	default:
		c.Fail(400, errors.New("unknown content-type: "+ctype))
		return false
	}
	return c.BindWith(obj, b)
}

func (c *Context) BindWith(obj interface{}, b binding.Binding) bool {
	if err := b.Bind(c.Request, obj); err != nil {
		c.Fail(400, err)
		return false
	}
	return true
}

func (c *Context) Render(code int, render render.Render, obj ...interface{}) {
	if err := render.Render(c.Writer, code, obj...); err != nil {
		c.ErrorTyped(err, ErrorTypeInternal, obj)
		c.Abort(500)
	}
}

// Serializes the given struct as JSON into the response body in a fast and efficient way.
// It also sets the Content-Type as "application/json".
func (c *Context) JSON(code int, obj interface{}) {
	c.Render(code, render.JSON, obj)
}

// Serializes the given struct as XML into the response body in a fast and efficient way.
// It also sets the Content-Type as "application/xml".
func (c *Context) XML(code int, obj interface{}) {
	c.Render(code, render.XML, obj)
}

// Renders the HTTP template specified by its file name.
// It also updates the HTTP code and sets the Content-Type as "text/html".
// See http://golang.org/doc/articles/wiki/
func (c *Context) HTML(code int, name string, obj interface{}) {
	c.Render(code, c.Engine.HTMLRender, name, obj)
}

// Writes the given string into the response body and sets the Content-Type to "text/plain".
func (c *Context) String(code int, format string, values ...interface{}) {
	c.Render(code, render.Plain, format, values)
}

// Writes some data into the body stream and updates the HTTP code.
func (c *Context) Data(code int, contentType string, data []byte) {
	if len(contentType) > 0 {
		c.Writer.Header().Set("Content-Type", contentType)
	}
	if code >= 0 {
		c.Writer.WriteHeader(code)
	}
	c.Writer.Write(data)
}

// Writes the specified file into the body stream
func (c *Context) File(filepath string) {
	http.ServeFile(c.Writer, c.Request, filepath)
}
<commit_msg>spelling<commit_after>package gin

import (
	"bytes"
	"errors"
	"fmt"
	"github.com/gin-gonic/gin/binding"
	"github.com/gin-gonic/gin/render"
	"github.com/julienschmidt/httprouter"
	"log"
	"net/http"
)

const (
	ErrorTypeInternal = 1 << iota
	ErrorTypeExternal = 1 << iota
	ErrorTypeAll      = 0xffffffff
)

// Used internally to collect errors that occurred during an http request.
type errorMsg struct {
	Err  string      `json:"error"`
	Type uint32      `json:"-"`
	Meta interface{} `json:"meta"`
}

type errorMsgs []errorMsg

func (a errorMsgs) ByType(typ uint32) errorMsgs {
	if len(a) == 0 {
		return a
	}
	result := make(errorMsgs, 0, len(a))
	for _, msg := range a {
		if msg.Type&typ > 0 {
			result = append(result, msg)
		}
	}
	return result
}

func (a errorMsgs) String() string {
	if len(a) == 0 {
		return ""
	}
	var buffer bytes.Buffer
	for i, msg := range a {
		text := fmt.Sprintf("Error #%02d: %s \n     Meta: %v\n", (i + 1), msg.Err, msg.Meta)
		buffer.WriteString(text)
	}
	return buffer.String()
}

// Context is the most important part of gin. It allows us to pass variables between middleware,
// manage the flow, validate the JSON of a request and render a JSON response for example.
type Context struct {
	writermem responseWriter
	Request   *http.Request
	Writer    ResponseWriter
	Keys      map[string]interface{}
	Errors    errorMsgs
	Params    httprouter.Params
	Engine    *Engine
	handlers  []HandlerFunc
	index     int8
}

/************************************/
/********** ROUTES GROUPING *********/
/************************************/

func (engine *Engine) createContext(w http.ResponseWriter, req *http.Request, params httprouter.Params, handlers []HandlerFunc) *Context {
	c := engine.cache.Get().(*Context)
	c.writermem.reset(w)
	c.Request = req
	c.Params = params
	c.handlers = handlers
	c.Keys = nil
	c.index = -1
	c.Errors = c.Errors[0:0]
	return c
}

/************************************/
/****** FLOW AND ERROR MANAGEMENT****/
/************************************/

func (c *Context) Copy() *Context {
	var cp Context = *c
	cp.index = AbortIndex
	cp.handlers = nil
	return &cp
}

// Next should be used only in the middlewares.
// It executes the pending handlers in the chain inside the calling handler.
// See example in github.
func (c *Context) Next() {
	c.index++
	s := int8(len(c.handlers))
	for ; c.index < s; c.index++ {
		c.handlers[c.index](c)
	}
}

// Forces the system to do not continue calling the pending handlers.
// For example, the first handler checks if the request is authorized. If it's not, context.Abort(401) should be called.
// The rest of pending handlers would never be called for that request.
func (c *Context) Abort(code int) {
	if code >= 0 {
		c.Writer.WriteHeader(code)
	}
	c.index = AbortIndex
}

// Fail is the same as Abort plus an error message.
// Calling `context.Fail(500, err)` is equivalent to:
// ```
// context.Error("Operation aborted", err)
// context.Abort(500)
// ```
func (c *Context) Fail(code int, err error) {
	c.Error(err, "Operation aborted")
	c.Abort(code)
}

func (c *Context) ErrorTyped(err error, typ uint32, meta interface{}) {
	c.Errors = append(c.Errors, errorMsg{
		Err:  err.Error(),
		Type: typ,
		Meta: meta,
	})
}

// Attaches an error to the current context. The error is pushed to a list of errors.
// It's a good idea to call Error for each error that occurred during the resolution of a request.
// A middleware can be used to collect all the errors and push them to a database together, print a log, or append it in the HTTP response.
func (c *Context) Error(err error, meta interface{}) {
	c.ErrorTyped(err, ErrorTypeExternal, meta)
}

func (c *Context) LastError() error {
	s := len(c.Errors)
	if s > 0 {
		return errors.New(c.Errors[s-1].Err)
	} else {
		return nil
	}
}

/************************************/
/******** METADATA MANAGEMENT********/
/************************************/

// Sets a new pair key/value just for the specified context.
// It also lazy initializes the hashmap.
func (c *Context) Set(key string, item interface{}) {
	if c.Keys == nil {
		c.Keys = make(map[string]interface{})
	}
	c.Keys[key] = item
}

// Get returns the value for the given key or an error if the key does not exist.
func (c *Context) Get(key string) (interface{}, error) {
	if c.Keys != nil {
		item, ok := c.Keys[key]
		if ok {
			return item, nil
		}
	}
	return nil, errors.New("Key does not exist.")
}

// MustGet returns the value for the given key or panics if the value doesn't exist.
func (c *Context) MustGet(key string) interface{} {
	value, err := c.Get(key)
	if err != nil || value == nil {
		log.Panicf("Key %s doesn't exist", key)
	}
	return value
}

/************************************/
/******** ENCODING MANAGEMENT********/
/************************************/

// This function checks the Content-Type to select a binding engine automatically,
// Depending the "Content-Type" header different bindings are used:
// "application/json" --> JSON binding
// "application/xml"  --> XML binding
// else --> returns an error
// if Parses the request's body as JSON if Content-Type == "application/json"  using JSON or XML  as a JSON input. It decodes the json payload into the struct specified as a pointer.Like ParseBody() but this method also writes a 400 error if the json is not valid.
func (c *Context) Bind(obj interface{}) bool {
	var b binding.Binding
	ctype := filterFlags(c.Request.Header.Get("Content-Type"))
	switch {
	case c.Request.Method == "GET" || ctype == MIMEPOSTForm:
		b = binding.Form
	case ctype == MIMEJSON:
		b = binding.JSON
	case ctype == MIMEXML || ctype == MIMEXML2:
		b = binding.XML
	default:
		c.Fail(400, errors.New("unknown content-type: "+ctype))
		return false
	}
	return c.BindWith(obj, b)
}

func (c *Context) BindWith(obj interface{}, b binding.Binding) bool {
	if err := b.Bind(c.Request, obj); err != nil {
		c.Fail(400, err)
		return false
	}
	return true
}

func (c *Context) Render(code int, render render.Render, obj ...interface{}) {
	if err := render.Render(c.Writer, code, obj...); err != nil {
		c.ErrorTyped(err, ErrorTypeInternal, obj)
		c.Abort(500)
	}
}

// Serializes the given struct as JSON into the response body in a fast and efficient way.
// It also sets the Content-Type as "application/json".
func (c *Context) JSON(code int, obj interface{}) {
	c.Render(code, render.JSON, obj)
}

// Serializes the given struct as XML into the response body in a fast and efficient way.
// It also sets the Content-Type as "application/xml".
func (c *Context) XML(code int, obj interface{}) {
	c.Render(code, render.XML, obj)
}

// Renders the HTTP template specified by its file name.
// It also updates the HTTP code and sets the Content-Type as "text/html".
// See http://golang.org/doc/articles/wiki/
func (c *Context) HTML(code int, name string, obj interface{}) {
	c.Render(code, c.Engine.HTMLRender, name, obj)
}

// Writes the given string into the response body and sets the Content-Type to "text/plain".
func (c *Context) String(code int, format string, values ...interface{}) {
	c.Render(code, render.Plain, format, values)
}

// Writes some data into the body stream and updates the HTTP code.
func (c *Context) Data(code int, contentType string, data []byte) {
	if len(contentType) > 0 {
		c.Writer.Header().Set("Content-Type", contentType)
	}
	if code >= 0 {
		c.Writer.WriteHeader(code)
	}
	c.Writer.Write(data)
}

// Writes the specified file into the body stream
func (c *Context) File(filepath string) {
	http.ServeFile(c.Writer, c.Request, filepath)
}
<|endoftext|>