<commit_before>/*
 * Copyright (c) 2013 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package goapp

import (
	"bytes"
	"code.google.com/p/go.net/html"
	"io"
	"net/url"
	"regexp"
	"strings"
)

func sanitizeLink(u *url.URL, v string) string {
	p, err := u.Parse(v)
	if err != nil {
		return ""
	}
	if !acceptableUriSchemes[p.Scheme] {
		return ""
	}

	return p.String()
}

func sanitizeStyle(v string) string {
	return v
}

func sanitizeAttributes(u *url.URL, t *html.Token) {
	var attrs []html.Attribute
	var isLink = false
	for _, a := range t.Attr {
		if a.Key == "target" {
		} else if a.Key == "style" {
			a.Val = sanitizeStyle(a.Val)
			attrs = append(attrs, a)
		} else if acceptableAttributes[a.Key] {
			if a.Key == "href" || a.Key == "src" {
				a.Val = sanitizeLink(u, a.Val)
			}
			if a.Key == "href" {
				isLink = true
			}
			attrs = append(attrs, a)
		}
	}
	if isLink {
		attrs = append(attrs, html.Attribute{
			Key: "target",
			Val: "_blank",
		})
	}
	t.Attr = attrs
}

func Sanitize(s string, u *url.URL) (string, string) {
	r := bytes.NewReader([]byte(s))
	z := html.NewTokenizer(r)
	buf := &bytes.Buffer{}
	snip := &bytes.Buffer{}
	skip := 0
	u.RawQuery = ""
	u.Fragment = ""
	for {
		if z.Next() == html.ErrorToken {
			if err := z.Err(); err == io.EOF {
				break
			} else {
				return s, snipper(s)
			}
		}

		t := z.Token()
		if t.Type == html.StartTagToken || t.Type == html.SelfClosingTagToken {
			if !acceptableElements[t.Data] {
				if unacceptableElementsWithEndTag[t.Data] && t.Type != html.SelfClosingTagToken {
					skip += 1
				}
			} else {
				sanitizeAttributes(u, &t)
				buf.WriteString(t.String())
			}
		} else if t.Type == html.EndTagToken {
			if !acceptableElements[t.Data] {
				if unacceptableElementsWithEndTag[t.Data] {
					skip -= 1
				}
			} else {
				buf.WriteString(t.String())
			}
		} else if skip == 0 {
			buf.WriteString(t.String())
			if t.Type == html.TextToken {
				snip.WriteString(t.String())
			}
		}
	}

	return buf.String(), snipper(snip.String())
}

const snipLen = 100

var snipRe = regexp.MustCompile("[\\s]+")

func snipper(s string) string {
	s = snipRe.ReplaceAllString(strings.TrimSpace(s), " ")
	s = html.UnescapeString(s)
	if len(s) <= snipLen {
		return s
	}
	s = s[:snipLen]
	i := strings.LastIndexAny(s, " .-!?")
	if i != -1 {
		return s[:i]
	}
	return cleanNonUTF8(s)
}

// Based on list from MDN's HTML5 element list
// https://developer.mozilla.org/en-US/docs/Web/Guide/HTML/HTML5/HTML5_element_list
var acceptableElements = map[string]bool{
	// Root element
	// "html": true,

	// Document metadata
	// "head":  true,
	// "title": true,
	// "base":  true,
	// "link":  true,
	// "meta":  true,
	// "style": true,

	// Scripting
	"noscript": true,
	// "script":   true,

	// Sections
	// "body":    true,
	"section": true,
	"nav":     true,
	"article": true,
	"aside":   true,
	"h1":      true,
	"h2":      true,
	"h3":      true,
	"h4":      true,
	"h5":      true,
	"h6":      true,
	"header":  true,
	"footer":  true,
	"address": true,
	"main":    true,

	// Grouping content
	"p":          true,
	"hr":         true,
	"pre":        true,
	"blockquote": true,
	"ol":         true,
	"ul":         true,
	"li":         true,
	"dl":         true,
	"dt":         true,
	"dd":         true,
	"figure":     true,
	"figcaption": true,
	"div":        true,

	// Text-level semantics
	"a":      true,
	"em":     true,
	"strong": true,
	"small":  true,
	"s":      true,
	"cite":   true,
	"q":      true,
	"dfn":    true,
	"abbr":   true,
	"data":   true,
	"time":   true,
	"code":   true,
	"var":    true,
	"samp":   true,
	"kbd":    true,
	"sub":    true,
	"sup":    true,
	"i":      true,
	"b":      true,
	"u":      true,
	"mark":   true,
	"ruby":   true,
	"rt":     true,
	"rp":     true,
	"bdi":    true,
	"bdo":    true,
	"span":   true,
	"br":     true,
	"wbr":    true,

	// Edits
	"ins": true,
	"del": true,

	// Embedded content
	"img": true,
	//"iframe":   true,
	//"embed":    true,
	"object": true,
	"param":  true,
	"video":  true,
	"audio":  true,
	"source": true,
	"track":  true,
	"canvas": true,
	"map":    true,
	"area":   true,
	"svg":    true,
	"math":   true,

	// Tabular data
	"table":    true,
	"caption":  true,
	"colgroup": true,
	"col":      true,
	"tbody":    true,
	"thead":    true,
	"tfoot":    true,
	"tr":       true,
	"td":       true,
	"th":       true,

	// Forms
	// "form":     true,
	// "fieldset": true,
	// "legend":   true,
	// "label":    true,
	// "input":    true,
	// "button":   true,
	// "select":   true,
	// "datalist": true,
	// "optgroup": true,
	// "option":   true,
	// "textarea": true,
	// "keygen":   true,
	// "output":   true,
	// "progress": true,
	// "meter":    true,

	// Interactive elements
	// "details":  true,
	// "summary":  true,
	// "menuitem": true,
	// "menu":     true,
}

var unacceptableElementsWithEndTag = map[string]bool{
	"script": true,
	"applet": true,
	"style":  true,
}

// Based on list from MDN's HTML attribute reference
// https://developer.mozilla.org/en-US/docs/Web/HTML/Attributes
var acceptableAttributes = map[string]bool{
	"accept":         true,
	"accept-charset": true,
	// "accesskey":       true,
	"action":       true,
	"align":        true,
	"alt":          true,
	"async":        true,
	"autocomplete": true,
	// "autofocus":       true,
	// "autoplay":        true,
	// "bgcolor":         true,
	// "border":          true,
	"buffered":        true,
	"challenge":       true,
	"charset":         true,
	"checked":         true,
	"cite":            true,
	// "class":           true,
	"code":            true,
	"codebase":        true,
	"color":           true,
	"cols":            true,
	"colspan":         true,
	"content":         true,
	"contenteditable": true,
	"contextmenu":     true,
	"controls":        true,
	"coords":          true,
	"data":            true,
	"data-custom":     true,
	"datetime":        true,
	"default":         true,
	"defer":           true,
	"dir":             true,
	"dirname":         true,
	"disabled":        true,
	"download":        true,
	"draggable":       true,
	"dropzone":        true,
	"enctype":         true,
	// "for":             true,
	"form":       true,
	"headers":    true,
	"height":     true,
	"hidden":     true,
	"high":       true,
	"href":       true,
	"hreflang":   true,
	"http-equiv": true,
	"icon":       true,
	// "id":              true,
	"ismap":       true,
	"itemprop":    true,
	"keytype":     true,
	"kind":        true,
	"label":       true,
	"lang":        true,
	"language":    true,
	"list":        true,
	"loop":        true,
	"low":         true,
	"manifest":    true,
	"max":         true,
	"maxlength":   true,
	"media":       true,
	"method":      true,
	"min":         true,
	"multiple":    true,
	"name":        true,
	"novalidate":  true,
	"open":        true,
	"optimum":     true,
	"pattern":     true,
	"ping":        true,
	"placeholder": true,
	"poster":      true,
	// "preload":         true,
	"pubdate":    true,
	"radiogroup": true,
	"readonly":   true,
	"rel":        true,
	"required":   true,
	"reversed":   true,
	"rows":       true,
	"rowspan":    true,
	"sandbox":    true,
	"spellcheck": true,
	"scope":      true,
	// "scoped":          true,
	// "seamless":        true,
	"selected": true,
	"shape":    true,
	"size":     true,
	"sizes":    true,
	"span":     true,
	"src":      true,
	// "srcdoc":          true,
	"srclang": true,
	"start":   true,
	// "step":            true,
	// "style":           true,
	"summary": true,
	// "tabindex":        true,
	// "target":          true,
	"title": true,
	"type":  true,
	// "usemap":          true,
	"value": true,
	"width": true,
	// "wrap":            true,
}

// Based on list from Wikipedia's URI scheme
// http://en.wikipedia.org/wiki/URI_scheme
var acceptableUriSchemes = map[string]bool{
	"aim":      true,
	"apt":      true,
	"bitcoin":  true,
	"callto":   true,
	"cvs":      true,
	"facetime": true,
	"feed":     true,
	"ftp":      true,
	"git":      true,
	"gopher":   true,
	"gtalk":    true,
	"http":     true,
	"https":    true,
	"imap":     true,
	"irc":      true,
	"itms":     true,
	"jabber":   true,
	"magnet":   true,
	"mailto":   true,
	"mms":      true,
	"msnim":    true,
	"news":     true,
	"nntp":     true,
	"rtmp":     true,
	"rtsp":     true,
	"sftp":     true,
	"skype":    true,
	"svn":      true,
	"ymsgr":    true,
}
<commit_msg>Trim incoming whitespace<commit_after>/*
 * Copyright (c) 2013 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package goapp

import (
	"bytes"
	"code.google.com/p/go.net/html"
	"io"
	"net/url"
	"regexp"
	"strings"
)

func sanitizeLink(u *url.URL, v string) string {
	p, err := u.Parse(v)
	if err != nil {
		return ""
	}
	if !acceptableUriSchemes[p.Scheme] {
		return ""
	}

	return p.String()
}

func sanitizeStyle(v string) string {
	return v
}

func sanitizeAttributes(u *url.URL, t *html.Token) {
	var attrs []html.Attribute
	var isLink = false
	for _, a := range t.Attr {
		if a.Key == "target" {
		} else if a.Key == "style" {
			a.Val = sanitizeStyle(a.Val)
			attrs = append(attrs, a)
		} else if acceptableAttributes[a.Key] {
			if a.Key == "href" || a.Key == "src" {
				a.Val = sanitizeLink(u, a.Val)
			}
			if a.Key == "href" {
				isLink = true
			}
			attrs = append(attrs, a)
		}
	}
	if isLink {
		attrs = append(attrs, html.Attribute{
			Key: "target",
			Val: "_blank",
		})
	}
	t.Attr = attrs
}

func Sanitize(s string, u *url.URL) (string, string) {
	r := bytes.NewReader([]byte(s))
	r := bytes.NewReader([]byte(strings.TrimSpace(s)))
	z := html.NewTokenizer(r)
	buf := &bytes.Buffer{}
	snip := &bytes.Buffer{}
	skip := 0
	u.RawQuery = ""
	u.Fragment = ""
	for {
		if z.Next() == html.ErrorToken {
			if err := z.Err(); err == io.EOF {
				break
			} else {
				return s, snipper(s)
			}
		}

		t := z.Token()
		if t.Type == html.StartTagToken || t.Type == html.SelfClosingTagToken {
			if !acceptableElements[t.Data] {
				if unacceptableElementsWithEndTag[t.Data] && t.Type != html.SelfClosingTagToken {
					skip += 1
				}
			} else {
				sanitizeAttributes(u, &t)
				buf.WriteString(t.String())
			}
		} else if t.Type == html.EndTagToken {
			if !acceptableElements[t.Data] {
				if unacceptableElementsWithEndTag[t.Data] {
					skip -= 1
				}
			} else {
				buf.WriteString(t.String())
			}
		} else if skip == 0 {
			buf.WriteString(t.String())
			if t.Type == html.TextToken {
				snip.WriteString(t.String())
			}
		}
	}

	return buf.String(), snipper(snip.String())
}

const snipLen = 100

var snipRe = regexp.MustCompile("[\\s]+")

func snipper(s string) string {
	s = snipRe.ReplaceAllString(strings.TrimSpace(s), " ")
	s = html.UnescapeString(s)
	if len(s) <= snipLen {
		return s
	}
	s = s[:snipLen]
	i := strings.LastIndexAny(s, " .-!?")
	if i != -1 {
		return s[:i]
	}
	return cleanNonUTF8(s)
}

// Based on list from MDN's HTML5 element list
// https://developer.mozilla.org/en-US/docs/Web/Guide/HTML/HTML5/HTML5_element_list
var acceptableElements = map[string]bool{
	// Root element
	// "html": true,

	// Document metadata
	// "head":  true,
	// "title": true,
	// "base":  true,
	// "link":  true,
	// "meta":  true,
	// "style": true,

	// Scripting
	"noscript": true,
	// "script":   true,

	// Sections
	// "body":    true,
	"section": true,
	"nav":     true,
	"article": true,
	"aside":   true,
	"h1":      true,
	"h2":      true,
	"h3":      true,
	"h4":      true,
	"h5":      true,
	"h6":      true,
	"header":  true,
	"footer":  true,
	"address": true,
	"main":    true,

	// Grouping content
	"p":          true,
	"hr":         true,
	"pre":        true,
	"blockquote": true,
	"ol":         true,
	"ul":         true,
	"li":         true,
	"dl":         true,
	"dt":         true,
	"dd":         true,
	"figure":     true,
	"figcaption": true,
	"div":        true,

	// Text-level semantics
	"a":      true,
	"em":     true,
	"strong": true,
	"small":  true,
	"s":      true,
	"cite":   true,
	"q":      true,
	"dfn":    true,
	"abbr":   true,
	"data":   true,
	"time":   true,
	"code":   true,
	"var":    true,
	"samp":   true,
	"kbd":    true,
	"sub":    true,
	"sup":    true,
	"i":      true,
	"b":      true,
	"u":      true,
	"mark":   true,
	"ruby":   true,
	"rt":     true,
	"rp":     true,
	"bdi":    true,
	"bdo":    true,
	"span":   true,
	"br":     true,
	"wbr":    true,

	// Edits
	"ins": true,
	"del": true,

	// Embedded content
	"img": true,
	//"iframe":   true,
	//"embed":    true,
	"object": true,
	"param":  true,
	"video":  true,
	"audio":  true,
	"source": true,
	"track":  true,
	"canvas": true,
	"map":    true,
	"area":   true,
	"svg":    true,
	"math":   true,

	// Tabular data
	"table":    true,
	"caption":  true,
	"colgroup": true,
	"col":      true,
	"tbody":    true,
	"thead":    true,
	"tfoot":    true,
	"tr":       true,
	"td":       true,
	"th":       true,

	// Forms
	// "form":     true,
	// "fieldset": true,
	// "legend":   true,
	// "label":    true,
	// "input":    true,
	// "button":   true,
	// "select":   true,
	// "datalist": true,
	// "optgroup": true,
	// "option":   true,
	// "textarea": true,
	// "keygen":   true,
	// "output":   true,
	// "progress": true,
	// "meter":    true,

	// Interactive elements
	// "details":  true,
	// "summary":  true,
	// "menuitem": true,
	// "menu":     true,
}

var unacceptableElementsWithEndTag = map[string]bool{
	"script": true,
	"applet": true,
	"style":  true,
}

// Based on list from MDN's HTML attribute reference
// https://developer.mozilla.org/en-US/docs/Web/HTML/Attributes
var acceptableAttributes = map[string]bool{
	"accept":         true,
	"accept-charset": true,
	// "accesskey":       true,
	"action":       true,
	"align":        true,
	"alt":          true,
	"async":        true,
	"autocomplete": true,
	// "autofocus":       true,
	// "autoplay":        true,
	// "bgcolor":         true,
	// "border":          true,
	"buffered":        true,
	"challenge":       true,
	"charset":         true,
	"checked":         true,
	"cite":            true,
	// "class":           true,
	"code":            true,
	"codebase":        true,
	"color":           true,
	"cols":            true,
	"colspan":         true,
	"content":         true,
	"contenteditable": true,
	"contextmenu":     true,
	"controls":        true,
	"coords":          true,
	"data":            true,
	"data-custom":     true,
	"datetime":        true,
	"default":         true,
	"defer":           true,
	"dir":             true,
	"dirname":         true,
	"disabled":        true,
	"download":        true,
	"draggable":       true,
	"dropzone":        true,
	"enctype":         true,
	// "for":             true,
	"form":       true,
	"headers":    true,
	"height":     true,
	"hidden":     true,
	"high":       true,
	"href":       true,
	"hreflang":   true,
	"http-equiv": true,
	"icon":       true,
	// "id":              true,
	"ismap":       true,
	"itemprop":    true,
	"keytype":     true,
	"kind":        true,
	"label":       true,
	"lang":        true,
	"language":    true,
	"list":        true,
	"loop":        true,
	"low":         true,
	"manifest":    true,
	"max":         true,
	"maxlength":   true,
	"media":       true,
	"method":      true,
	"min":         true,
	"multiple":    true,
	"name":        true,
	"novalidate":  true,
	"open":        true,
	"optimum":     true,
	"pattern":     true,
	"ping":        true,
	"placeholder": true,
	"poster":      true,
	// "preload":         true,
	"pubdate":    true,
	"radiogroup": true,
	"readonly":   true,
	"rel":        true,
	"required":   true,
	"reversed":   true,
	"rows":       true,
	"rowspan":    true,
	"sandbox":    true,
	"spellcheck": true,
	"scope":      true,
	// "scoped":          true,
	// "seamless":        true,
	"selected": true,
	"shape":    true,
	"size":     true,
	"sizes":    true,
	"span":     true,
	"src":      true,
	// "srcdoc":          true,
	"srclang": true,
	"start":   true,
	// "step":            true,
	// "style":           true,
	"summary": true,
	// "tabindex":        true,
	// "target":          true,
	"title": true,
	"type":  true,
	// "usemap":          true,
	"value": true,
	"width": true,
	// "wrap":            true,
}

// Based on list from Wikipedia's URI scheme
// http://en.wikipedia.org/wiki/URI_scheme
var acceptableUriSchemes = map[string]bool{
	"aim":      true,
	"apt":      true,
	"bitcoin":  true,
	"callto":   true,
	"cvs":      true,
	"facetime": true,
	"feed":     true,
	"ftp":      true,
	"git":      true,
	"gopher":   true,
	"gtalk":    true,
	"http":     true,
	"https":    true,
	"imap":     true,
	"irc":      true,
	"itms":     true,
	"jabber":   true,
	"magnet":   true,
	"mailto":   true,
	"mms":      true,
	"msnim":    true,
	"news":     true,
	"nntp":     true,
	"rtmp":     true,
	"rtsp":     true,
	"sftp":     true,
	"skype":    true,
	"svn":      true,
	"ymsgr":    true,
}
<|endoftext|>


<commit_before>// Copyright 2020 The Matrix.org Foundation C.I.C.
//
// 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 perform

import (
	"context"
	"fmt"

	federationSenderAPI "github.com/matrix-org/dendrite/federationsender/api"
	"github.com/matrix-org/dendrite/roomserver/api"
	"github.com/matrix-org/dendrite/roomserver/internal/helpers"
	"github.com/matrix-org/dendrite/roomserver/internal/input"
	"github.com/matrix-org/dendrite/roomserver/state"
	"github.com/matrix-org/dendrite/roomserver/storage"
	"github.com/matrix-org/dendrite/roomserver/types"
	"github.com/matrix-org/dendrite/setup/config"
	"github.com/matrix-org/gomatrixserverlib"
	log "github.com/sirupsen/logrus"
)

type Inviter struct {
	DB      storage.Database
	Cfg     *config.RoomServer
	FSAPI   federationSenderAPI.FederationSenderInternalAPI
	Inputer *input.Inputer
}

// nolint:gocyclo
func (r *Inviter) PerformInvite(
	ctx context.Context,
	req *api.PerformInviteRequest,
	res *api.PerformInviteResponse,
) ([]api.OutputEvent, error) {
	event := req.Event
	if event.StateKey() == nil {
		return nil, fmt.Errorf("invite must be a state event")
	}

	roomID := event.RoomID()
	targetUserID := *event.StateKey()
	info, err := r.DB.RoomInfo(ctx, roomID)
	if err != nil {
		return nil, fmt.Errorf("Failed to load RoomInfo: %w", err)
	}

	log.WithFields(log.Fields{
		"event_id":         event.EventID(),
		"room_id":          roomID,
		"room_version":     req.RoomVersion,
		"target_user_id":   targetUserID,
		"room_info_exists": info != nil,
	}).Info("processing invite event")

	_, domain, _ := gomatrixserverlib.SplitID('@', targetUserID)
	isTargetLocal := domain == r.Cfg.Matrix.ServerName
	isOriginLocal := event.Origin() == r.Cfg.Matrix.ServerName

	inviteState := req.InviteRoomState
	if len(inviteState) == 0 && info != nil {
		var is []gomatrixserverlib.InviteV2StrippedState
		if is, err = buildInviteStrippedState(ctx, r.DB, info, req); err == nil {
			inviteState = is
		}
	}
	if len(inviteState) == 0 {
		if err = event.SetUnsignedField("invite_room_state", struct{}{}); err != nil {
			return nil, fmt.Errorf("event.SetUnsignedField: %w", err)
		}
	} else {
		if err = event.SetUnsignedField("invite_room_state", inviteState); err != nil {
			return nil, fmt.Errorf("event.SetUnsignedField: %w", err)
		}
	}

	var isAlreadyJoined bool
	if info != nil {
		_, isAlreadyJoined, _, err = r.DB.GetMembership(ctx, info.RoomNID, *event.StateKey())
		if err != nil {
			return nil, fmt.Errorf("r.DB.GetMembership: %w", err)
		}
	}
	if isAlreadyJoined {
		// If the user is joined to the room then that takes precedence over this
		// invite event. It makes little sense to move a user that is already
		// joined to the room into the invite state.
		// This could plausibly happen if an invite request raced with a join
		// request for a user. For example if a user was invited to a public
		// room and they joined the room at the same time as the invite was sent.
		// The other way this could plausibly happen is if an invite raced with
		// a kick. For example if a user was kicked from a room in error and in
		// response someone else in the room re-invited them then it is possible
		// for the invite request to race with the leave event so that the
		// target receives invite before it learns that it has been kicked.
		// There are a few ways this could be plausibly handled in the roomserver.
		// 1) Store the invite, but mark it as retired. That will result in the
		//    permanent rejection of that invite event. So even if the target
		//    user leaves the room and the invite is retransmitted it will be
		//    ignored. However a new invite with a new event ID would still be
		//    accepted.
		// 2) Silently discard the invite event. This means that if the event
		//    was retransmitted at a later date after the target user had left
		//    the room we would accept the invite. However since we hadn't told
		//    the sending server that the invite had been discarded it would
		//    have no reason to attempt to retry.
		// 3) Signal the sending server that the user is already joined to the
		//    room.
		// For now we will implement option 2. Since in the abesence of a retry
		// mechanism it will be equivalent to option 1, and we don't have a
		// signalling mechanism to implement option 3.
		res.Error = &api.PerformError{
			Code: api.PerformErrorNotAllowed,
			Msg:  "User is already joined to room",
		}
		return nil, nil
	}

	if isOriginLocal {
		// The invite originated locally. Therefore we have a responsibility to
		// try and see if the user is allowed to make this invite. We can't do
		// this for invites coming in over federation - we have to take those on
		// trust.
		_, err = helpers.CheckAuthEvents(ctx, r.DB, event, event.AuthEventIDs())
		if err != nil {
			log.WithError(err).WithField("event_id", event.EventID()).WithField("auth_event_ids", event.AuthEventIDs()).Error(
				"processInviteEvent.checkAuthEvents failed for event",
			)
			res.Error = &api.PerformError{
				Msg:  err.Error(),
				Code: api.PerformErrorNotAllowed,
			}
		}

		// If the invite originated from us and the target isn't local then we
		// should try and send the invite over federation first. It might be
		// that the remote user doesn't exist, in which case we can give up
		// processing here.
		if req.SendAsServer != api.DoNotSendToOtherServers && !isTargetLocal {
			fsReq := &federationSenderAPI.PerformInviteRequest{
				RoomVersion:     req.RoomVersion,
				Event:           event,
				InviteRoomState: inviteState,
			}
			fsRes := &federationSenderAPI.PerformInviteResponse{}
			if err = r.FSAPI.PerformInvite(ctx, fsReq, fsRes); err != nil {
				res.Error = &api.PerformError{
					Msg:  err.Error(),
					Code: api.PerformErrorNotAllowed,
				}
				log.WithError(err).WithField("event_id", event.EventID()).Error("r.FSAPI.PerformInvite failed")
				return nil, nil
			}
			event = fsRes.Event
		}

		// Send the invite event to the roomserver input stream. This will
		// notify existing users in the room about the invite, update the
		// membership table and ensure that the event is ready and available
		// to use as an auth event when accepting the invite.
		inputReq := &api.InputRoomEventsRequest{
			InputRoomEvents: []api.InputRoomEvent{
				{
					Kind:         api.KindNew,
					Event:        event,
					AuthEventIDs: event.AuthEventIDs(),
					SendAsServer: req.SendAsServer,
				},
			},
		}
		inputRes := &api.InputRoomEventsResponse{}
		r.Inputer.InputRoomEvents(context.Background(), inputReq, inputRes)
		if err = inputRes.Err(); err != nil {
			res.Error = &api.PerformError{
				Msg:  fmt.Sprintf("r.InputRoomEvents: %s", err.Error()),
				Code: api.PerformErrorNotAllowed,
			}
			log.WithError(err).WithField("event_id", event.EventID()).Error("r.InputRoomEvents failed")
			return nil, nil
		}
	} else {
		// The invite originated over federation. Process the membership
		// update, which will notify the sync API etc about the incoming
		// invite.
		updater, err := r.DB.MembershipUpdater(ctx, roomID, targetUserID, isTargetLocal, req.RoomVersion)
		if err != nil {
			return nil, fmt.Errorf("r.DB.MembershipUpdater: %w", err)
		}

		unwrapped := event.Unwrap()
		outputUpdates, err := helpers.UpdateToInviteMembership(updater, unwrapped, nil, req.Event.RoomVersion)
		if err != nil {
			return nil, fmt.Errorf("updateToInviteMembership: %w", err)
		}

		if err = updater.Commit(); err != nil {
			return nil, fmt.Errorf("updater.Commit: %w", err)
		}

		return outputUpdates, nil
	}

	return nil, nil
}

func buildInviteStrippedState(
	ctx context.Context,
	db storage.Database,
	info *types.RoomInfo,
	input *api.PerformInviteRequest,
) ([]gomatrixserverlib.InviteV2StrippedState, error) {
	stateWanted := []gomatrixserverlib.StateKeyTuple{}
	// "If they are set on the room, at least the state for m.room.avatar, m.room.canonical_alias, m.room.join_rules, and m.room.name SHOULD be included."
	// https://matrix.org/docs/spec/client_server/r0.6.0#m-room-member
	for _, t := range []string{
		gomatrixserverlib.MRoomName, gomatrixserverlib.MRoomCanonicalAlias,
		gomatrixserverlib.MRoomAliases, gomatrixserverlib.MRoomJoinRules,
		"m.room.avatar", "m.room.encryption",
	} {
		stateWanted = append(stateWanted, gomatrixserverlib.StateKeyTuple{
			EventType: t,
			StateKey:  "",
		})
	}
	roomState := state.NewStateResolution(db, *info)
	stateEntries, err := roomState.LoadStateAtSnapshotForStringTuples(
		ctx, info.StateSnapshotNID, stateWanted,
	)
	if err != nil {
		return nil, err
	}
	stateNIDs := []types.EventNID{}
	for _, stateNID := range stateEntries {
		stateNIDs = append(stateNIDs, stateNID.EventNID)
	}
	stateEvents, err := db.Events(ctx, stateNIDs)
	if err != nil {
		return nil, err
	}
	inviteState := []gomatrixserverlib.InviteV2StrippedState{
		gomatrixserverlib.NewInviteV2StrippedState(input.Event.Event),
	}
	stateEvents = append(stateEvents, types.Event{Event: input.Event.Unwrap()})
	for _, event := range stateEvents {
		inviteState = append(inviteState, gomatrixserverlib.NewInviteV2StrippedState(event.Event))
	}
	return inviteState, nil
}
<commit_msg>Add m.room.create to invite stripped state (#1740)<commit_after>// Copyright 2020 The Matrix.org Foundation C.I.C.
//
// 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 perform

import (
	"context"
	"fmt"

	federationSenderAPI "github.com/matrix-org/dendrite/federationsender/api"
	"github.com/matrix-org/dendrite/roomserver/api"
	"github.com/matrix-org/dendrite/roomserver/internal/helpers"
	"github.com/matrix-org/dendrite/roomserver/internal/input"
	"github.com/matrix-org/dendrite/roomserver/state"
	"github.com/matrix-org/dendrite/roomserver/storage"
	"github.com/matrix-org/dendrite/roomserver/types"
	"github.com/matrix-org/dendrite/setup/config"
	"github.com/matrix-org/gomatrixserverlib"
	log "github.com/sirupsen/logrus"
)

type Inviter struct {
	DB      storage.Database
	Cfg     *config.RoomServer
	FSAPI   federationSenderAPI.FederationSenderInternalAPI
	Inputer *input.Inputer
}

// nolint:gocyclo
func (r *Inviter) PerformInvite(
	ctx context.Context,
	req *api.PerformInviteRequest,
	res *api.PerformInviteResponse,
) ([]api.OutputEvent, error) {
	event := req.Event
	if event.StateKey() == nil {
		return nil, fmt.Errorf("invite must be a state event")
	}

	roomID := event.RoomID()
	targetUserID := *event.StateKey()
	info, err := r.DB.RoomInfo(ctx, roomID)
	if err != nil {
		return nil, fmt.Errorf("Failed to load RoomInfo: %w", err)
	}

	log.WithFields(log.Fields{
		"event_id":         event.EventID(),
		"room_id":          roomID,
		"room_version":     req.RoomVersion,
		"target_user_id":   targetUserID,
		"room_info_exists": info != nil,
	}).Info("processing invite event")

	_, domain, _ := gomatrixserverlib.SplitID('@', targetUserID)
	isTargetLocal := domain == r.Cfg.Matrix.ServerName
	isOriginLocal := event.Origin() == r.Cfg.Matrix.ServerName

	inviteState := req.InviteRoomState
	if len(inviteState) == 0 && info != nil {
		var is []gomatrixserverlib.InviteV2StrippedState
		if is, err = buildInviteStrippedState(ctx, r.DB, info, req); err == nil {
			inviteState = is
		}
	}
	if len(inviteState) == 0 {
		if err = event.SetUnsignedField("invite_room_state", struct{}{}); err != nil {
			return nil, fmt.Errorf("event.SetUnsignedField: %w", err)
		}
	} else {
		if err = event.SetUnsignedField("invite_room_state", inviteState); err != nil {
			return nil, fmt.Errorf("event.SetUnsignedField: %w", err)
		}
	}

	var isAlreadyJoined bool
	if info != nil {
		_, isAlreadyJoined, _, err = r.DB.GetMembership(ctx, info.RoomNID, *event.StateKey())
		if err != nil {
			return nil, fmt.Errorf("r.DB.GetMembership: %w", err)
		}
	}
	if isAlreadyJoined {
		// If the user is joined to the room then that takes precedence over this
		// invite event. It makes little sense to move a user that is already
		// joined to the room into the invite state.
		// This could plausibly happen if an invite request raced with a join
		// request for a user. For example if a user was invited to a public
		// room and they joined the room at the same time as the invite was sent.
		// The other way this could plausibly happen is if an invite raced with
		// a kick. For example if a user was kicked from a room in error and in
		// response someone else in the room re-invited them then it is possible
		// for the invite request to race with the leave event so that the
		// target receives invite before it learns that it has been kicked.
		// There are a few ways this could be plausibly handled in the roomserver.
		// 1) Store the invite, but mark it as retired. That will result in the
		//    permanent rejection of that invite event. So even if the target
		//    user leaves the room and the invite is retransmitted it will be
		//    ignored. However a new invite with a new event ID would still be
		//    accepted.
		// 2) Silently discard the invite event. This means that if the event
		//    was retransmitted at a later date after the target user had left
		//    the room we would accept the invite. However since we hadn't told
		//    the sending server that the invite had been discarded it would
		//    have no reason to attempt to retry.
		// 3) Signal the sending server that the user is already joined to the
		//    room.
		// For now we will implement option 2. Since in the abesence of a retry
		// mechanism it will be equivalent to option 1, and we don't have a
		// signalling mechanism to implement option 3.
		res.Error = &api.PerformError{
			Code: api.PerformErrorNotAllowed,
			Msg:  "User is already joined to room",
		}
		return nil, nil
	}

	if isOriginLocal {
		// The invite originated locally. Therefore we have a responsibility to
		// try and see if the user is allowed to make this invite. We can't do
		// this for invites coming in over federation - we have to take those on
		// trust.
		_, err = helpers.CheckAuthEvents(ctx, r.DB, event, event.AuthEventIDs())
		if err != nil {
			log.WithError(err).WithField("event_id", event.EventID()).WithField("auth_event_ids", event.AuthEventIDs()).Error(
				"processInviteEvent.checkAuthEvents failed for event",
			)
			res.Error = &api.PerformError{
				Msg:  err.Error(),
				Code: api.PerformErrorNotAllowed,
			}
		}

		// If the invite originated from us and the target isn't local then we
		// should try and send the invite over federation first. It might be
		// that the remote user doesn't exist, in which case we can give up
		// processing here.
		if req.SendAsServer != api.DoNotSendToOtherServers && !isTargetLocal {
			fsReq := &federationSenderAPI.PerformInviteRequest{
				RoomVersion:     req.RoomVersion,
				Event:           event,
				InviteRoomState: inviteState,
			}
			fsRes := &federationSenderAPI.PerformInviteResponse{}
			if err = r.FSAPI.PerformInvite(ctx, fsReq, fsRes); err != nil {
				res.Error = &api.PerformError{
					Msg:  err.Error(),
					Code: api.PerformErrorNotAllowed,
				}
				log.WithError(err).WithField("event_id", event.EventID()).Error("r.FSAPI.PerformInvite failed")
				return nil, nil
			}
			event = fsRes.Event
		}

		// Send the invite event to the roomserver input stream. This will
		// notify existing users in the room about the invite, update the
		// membership table and ensure that the event is ready and available
		// to use as an auth event when accepting the invite.
		inputReq := &api.InputRoomEventsRequest{
			InputRoomEvents: []api.InputRoomEvent{
				{
					Kind:         api.KindNew,
					Event:        event,
					AuthEventIDs: event.AuthEventIDs(),
					SendAsServer: req.SendAsServer,
				},
			},
		}
		inputRes := &api.InputRoomEventsResponse{}
		r.Inputer.InputRoomEvents(context.Background(), inputReq, inputRes)
		if err = inputRes.Err(); err != nil {
			res.Error = &api.PerformError{
				Msg:  fmt.Sprintf("r.InputRoomEvents: %s", err.Error()),
				Code: api.PerformErrorNotAllowed,
			}
			log.WithError(err).WithField("event_id", event.EventID()).Error("r.InputRoomEvents failed")
			return nil, nil
		}
	} else {
		// The invite originated over federation. Process the membership
		// update, which will notify the sync API etc about the incoming
		// invite.
		updater, err := r.DB.MembershipUpdater(ctx, roomID, targetUserID, isTargetLocal, req.RoomVersion)
		if err != nil {
			return nil, fmt.Errorf("r.DB.MembershipUpdater: %w", err)
		}

		unwrapped := event.Unwrap()
		outputUpdates, err := helpers.UpdateToInviteMembership(updater, unwrapped, nil, req.Event.RoomVersion)
		if err != nil {
			return nil, fmt.Errorf("updateToInviteMembership: %w", err)
		}

		if err = updater.Commit(); err != nil {
			return nil, fmt.Errorf("updater.Commit: %w", err)
		}

		return outputUpdates, nil
	}

	return nil, nil
}

func buildInviteStrippedState(
	ctx context.Context,
	db storage.Database,
	info *types.RoomInfo,
	input *api.PerformInviteRequest,
) ([]gomatrixserverlib.InviteV2StrippedState, error) {
	stateWanted := []gomatrixserverlib.StateKeyTuple{}
	// "If they are set on the room, at least the state for m.room.avatar, m.room.canonical_alias, m.room.join_rules, and m.room.name SHOULD be included."
	// https://matrix.org/docs/spec/client_server/r0.6.0#m-room-member
	for _, t := range []string{
		gomatrixserverlib.MRoomName, gomatrixserverlib.MRoomCanonicalAlias,
		gomatrixserverlib.MRoomAliases, gomatrixserverlib.MRoomJoinRules,
		"m.room.avatar", "m.room.encryption", gomatrixserverlib.MRoomCreate,
	} {
		stateWanted = append(stateWanted, gomatrixserverlib.StateKeyTuple{
			EventType: t,
			StateKey:  "",
		})
	}
	roomState := state.NewStateResolution(db, *info)
	stateEntries, err := roomState.LoadStateAtSnapshotForStringTuples(
		ctx, info.StateSnapshotNID, stateWanted,
	)
	if err != nil {
		return nil, err
	}
	stateNIDs := []types.EventNID{}
	for _, stateNID := range stateEntries {
		stateNIDs = append(stateNIDs, stateNID.EventNID)
	}
	stateEvents, err := db.Events(ctx, stateNIDs)
	if err != nil {
		return nil, err
	}
	inviteState := []gomatrixserverlib.InviteV2StrippedState{
		gomatrixserverlib.NewInviteV2StrippedState(input.Event.Event),
	}
	stateEvents = append(stateEvents, types.Event{Event: input.Event.Unwrap()})
	for _, event := range stateEvents {
		inviteState = append(inviteState, gomatrixserverlib.NewInviteV2StrippedState(event.Event))
	}
	return inviteState, nil
}
<|endoftext|>


<commit_before>// Support for SPI-based RGB LED chains (APA-102).
//
// The first LED is the status LED:
//	blue: idle
//  green: connected
//  red: errors
//  orange: connected+errors
//
// The remaining LEDs are tally LEDs, using the sequential ID numbering
//  blue: found
//  green: preview
//  red: program
//  orange: program+preview
package spiled

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"github.com/kidoman/embd"
	_ "github.com/kidoman/embd/host/rpi" // This loads the RPi driver
	"github.com/qmsk/e2/tally"
	"log"
	"strings"
	"sync"
	"time"
)

type Protocol string

const APA102	Protocol	= "apa102"
const APA102X				= "apa102x"

type Options struct {
	Channel		byte	`long:"spiled-channel" metavar:"N" description:"/dev/spidev0.N"`
	Speed		int		`long:"spiled-speed" metavar:"HZ"`
	Protocol	string	`long:"spiled-protocol" metavar:"apa102|apa102x" description:"Type of LED"`
	Count		uint	`long:"spiled-count" metavar:"COUNT" description:"Number of LEDs"`
	Debug		bool	`long:"spiled-debug" description:"Dump SPI output"`
	Intensity	uint8	`long:"spiled-intensity" metavar:"0-255" default:"255"`
	Refresh		float64 `long:"spiled-refresh" metavar:"HZ" default:"10"`

	TallyIdle		LED		`long:"spiled-tally-idle"    metavar:"RRGGBB" default:"000040"`
	TallyPreview	LED		`long:"spiled-tally-preview" metavar:"RRGGBB" default:"00ff00"`
	TallyProgram	LED		`long:"spiled-tally-program" metavar:"RRGGBB" default:"ff0000"`
	TallyBoth		LED		`long:"spiled-tally-both"    metavar:"RRGGBB" default:"ff4000"`

	StatusIdle		LED		`long:"spiled-status-idle"    metavar:"RRGGBB" default:"0000ff"`
	StatusOK		LED		`long:"spiled-status-ok"      metavar:"RRGGBB" default:"00ff00"`
	StatusWarn	    LED		`long:"spiled-status-warn"    metavar:"RRGGBB" default:"ffff00"`
	StatusError		LED		`long:"spiled-status-error"   metavar:"RRGGBB" default:"ff0000"`
}

func (options Options) Make() (*SPILED, error) {
	var spiled = SPILED{
		options:   options,
	}

	if err := spiled.init(options); err != nil {
		return nil, err
	}

	return &spiled, nil
}

type SPILED struct {
	options     Options
	protocol	Protocol
	count		uint

	spiBus	embd.SPIBus

	leds		[]LED

	tallyChan chan tally.State
	waitGroup sync.WaitGroup
}

func (spiled *SPILED) init(options Options) error {
	if err := embd.InitSPI(); err != nil {
		return fmt.Errorf("embd.InitSPI: %v", err)
	}

	// SPI
	var spiMode byte = embd.SPIMode0
	var spiChannel byte = options.Channel // /dev/spidev0.X
	var spiSpeed int = options.Speed // Hz
	var spiBitsPerWord int = 8 // bits
	var spiDelay int = 0 // us?

	spiled.protocol = Protocol(strings.ToLower(options.Protocol))
	spiled.count = options.Count

	switch spiled.protocol {
	case APA102, APA102X:
		spiMode = embd.SPIMode0
		spiBitsPerWord = 8

	default:
		return fmt.Errorf("Invalid --spiled-protocol=%v", options.Protocol)
	}

	spiled.spiBus = embd.NewSPIBus(spiMode, spiChannel, spiSpeed, spiBitsPerWord, spiDelay)

	// initial output
	spiled.leds = make([]LED, spiled.count)

	if err := spiled.write(spiled.leds, time.Time{}); err != nil {
		return err
	}

	log.Printf("SPI-LED: Open %v with %d %s LEDs", spiled.spiBus, options.Count, spiled.protocol)

	return nil
}

func (spiled *SPILED) write(leds []LED, renderTime time.Time) error {
	var packet bytes.Buffer

	var stopByte = []byte{0xff}
	var stopCount = 4 * (1 + len(leds) / 32) // one bit per byte, in frames of 32 bits

	switch spiled.protocol {
	case APA102X:
		// variation where the stop frame must be 0x00
		stopByte = []byte{0x00}
	}

	// start
	var startFrame = []byte{0x00, 0x00, 0x00, 0x00}

	packet.Write(startFrame)

	// data
	for _, led := range leds {
		var ledFrame = led.render(renderTime)

		packet.Write(ledFrame)
	}

	// stop
	var stopFrame = bytes.Repeat(stopByte, stopCount)

	packet.Write(stopFrame)

	if spiled.options.Debug {
		fmt.Println(hex.Dump(packet.Bytes()))
	}

	// send
	if _, err := spiled.spiBus.Write(packet.Bytes()); err != nil {
		return fmt.Errorf("SPIBUs.Write %v: %v", packet.Len(), err)
	} else {
		return nil
	}
}

func (spiled *SPILED) close() {
	defer spiled.waitGroup.Done()

	log.Printf("SPI-LED: Close...")

	if err := spiled.spiBus.Close(); err != nil {
		log.Printf("embd.SPIBus.Close: %v", err)
	}
}

func (spiled *SPILED) updateTally(tallyState tally.State) {
	log.Printf("SPI-LED: Update tally State:")

	leds := make([]LED, spiled.count)

	var found, errors int

	for i, led := range leds {
		if i == 0 {
			// skip status
			continue
		}

		id := tally.ID(i)

		if tally, exists := tallyState.Tally[id]; !exists {
			// missing tally state for pin
		} else {
			found++

			if tally.Status.Program && tally.Status.Preview {
				led = spiled.options.TallyBoth
			} else if tally.Status.Preview {
				led = spiled.options.TallyPreview
			} else if tally.Status.Program {
				led = spiled.options.TallyProgram
			} else {
				led = spiled.options.TallyIdle
			}

			led.Intensity = spiled.options.Intensity

			if tally.Errors != nil {
				led.Strobe(1 * time.Second)
			}

			log.Printf("SPI-LED %v: id=%v status=%v errors=%v led=%v", i, id, tally.Status, len(tally.Errors), led)

		}

		leds[i] = led
	}

	errors = len(tallyState.Errors)

	// status LED
	var statusLED LED

	if found > 0 && errors > 0 {
		statusLED = spiled.options.StatusWarn
	} else if errors > 0 {
		statusLED = spiled.options.StatusError
	} else if found > 0 {
		statusLED = spiled.options.StatusOK
	} else {
		statusLED = spiled.options.StatusIdle
	}

	statusLED.Intensity = spiled.options.Intensity

	log.Printf("SPI-LED: found=%v errors=%v led=%v", found, errors, statusLED)

	leds[0] = statusLED

	// refresh
	spiled.leds = leds
}

func (spiled *SPILED) run() {
	defer spiled.close()

	refreshTimer := time.Tick(time.Duration(1.0 / spiled.options.Refresh * float64(time.Second)))

	for {
		select {
		case tallyState, ok := <-spiled.tallyChan:
			if ok {
				spiled.updateTally(tallyState)
			} else {
				return
			}
		case refreshTime := <-refreshTimer:
			if err := spiled.write(spiled.leds, refreshTime); err != nil {
				log.Printf("SPI-LED: Write error: %v", err)
			}
		}
	}

	log.Printf("SPI-LED: Done")
}

func (spiled *SPILED) RegisterTally(t *tally.Tally) {
	spiled.tallyChan = make(chan tally.State)
	spiled.waitGroup.Add(1)

	go spiled.run()

	t.Register(spiled.tallyChan)
}

// Close and Wait..
func (spiled *SPILED) Close() {
	log.Printf("SPI-LED: Close..")

	if spiled.tallyChan != nil {
		close(spiled.tallyChan)
	}

	spiled.waitGroup.Wait()
}
<commit_msg>spiled: even dimmer tally-idle<commit_after>// Support for SPI-based RGB LED chains (APA-102).
//
// The first LED is the status LED:
//	blue: idle
//  green: connected
//  red: errors
//  orange: connected+errors
//
// The remaining LEDs are tally LEDs, using the sequential ID numbering
//  blue: found
//  green: preview
//  red: program
//  orange: program+preview
package spiled

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"github.com/kidoman/embd"
	_ "github.com/kidoman/embd/host/rpi" // This loads the RPi driver
	"github.com/qmsk/e2/tally"
	"log"
	"strings"
	"sync"
	"time"
)

type Protocol string

const APA102	Protocol	= "apa102"
const APA102X				= "apa102x"

type Options struct {
	Channel		byte	`long:"spiled-channel" metavar:"N" description:"/dev/spidev0.N"`
	Speed		int		`long:"spiled-speed" metavar:"HZ"`
	Protocol	string	`long:"spiled-protocol" metavar:"apa102|apa102x" description:"Type of LED"`
	Count		uint	`long:"spiled-count" metavar:"COUNT" description:"Number of LEDs"`
	Debug		bool	`long:"spiled-debug" description:"Dump SPI output"`
	Intensity	uint8	`long:"spiled-intensity" metavar:"0-255" default:"255"`
	Refresh		float64 `long:"spiled-refresh" metavar:"HZ" default:"10"`

	TallyIdle		LED		`long:"spiled-tally-idle"    metavar:"RRGGBB" default:"000010"`
	TallyPreview	LED		`long:"spiled-tally-preview" metavar:"RRGGBB" default:"00ff00"`
	TallyProgram	LED		`long:"spiled-tally-program" metavar:"RRGGBB" default:"ff0000"`
	TallyBoth		LED		`long:"spiled-tally-both"    metavar:"RRGGBB" default:"ff4000"`

	StatusIdle		LED		`long:"spiled-status-idle"    metavar:"RRGGBB" default:"0000ff"`
	StatusOK		LED		`long:"spiled-status-ok"      metavar:"RRGGBB" default:"00ff00"`
	StatusWarn	    LED		`long:"spiled-status-warn"    metavar:"RRGGBB" default:"ffff00"`
	StatusError		LED		`long:"spiled-status-error"   metavar:"RRGGBB" default:"ff0000"`
}

func (options Options) Make() (*SPILED, error) {
	var spiled = SPILED{
		options:   options,
	}

	if err := spiled.init(options); err != nil {
		return nil, err
	}

	return &spiled, nil
}

type SPILED struct {
	options     Options
	protocol	Protocol
	count		uint

	spiBus	embd.SPIBus

	leds		[]LED

	tallyChan chan tally.State
	waitGroup sync.WaitGroup
}

func (spiled *SPILED) init(options Options) error {
	if err := embd.InitSPI(); err != nil {
		return fmt.Errorf("embd.InitSPI: %v", err)
	}

	// SPI
	var spiMode byte = embd.SPIMode0
	var spiChannel byte = options.Channel // /dev/spidev0.X
	var spiSpeed int = options.Speed // Hz
	var spiBitsPerWord int = 8 // bits
	var spiDelay int = 0 // us?

	spiled.protocol = Protocol(strings.ToLower(options.Protocol))
	spiled.count = options.Count

	switch spiled.protocol {
	case APA102, APA102X:
		spiMode = embd.SPIMode0
		spiBitsPerWord = 8

	default:
		return fmt.Errorf("Invalid --spiled-protocol=%v", options.Protocol)
	}

	spiled.spiBus = embd.NewSPIBus(spiMode, spiChannel, spiSpeed, spiBitsPerWord, spiDelay)

	// initial output
	spiled.leds = make([]LED, spiled.count)

	if err := spiled.write(spiled.leds, time.Time{}); err != nil {
		return err
	}

	log.Printf("SPI-LED: Open %v with %d %s LEDs", spiled.spiBus, options.Count, spiled.protocol)

	return nil
}

func (spiled *SPILED) write(leds []LED, renderTime time.Time) error {
	var packet bytes.Buffer

	var stopByte = []byte{0xff}
	var stopCount = 4 * (1 + len(leds) / 32) // one bit per byte, in frames of 32 bits

	switch spiled.protocol {
	case APA102X:
		// variation where the stop frame must be 0x00
		stopByte = []byte{0x00}
	}

	// start
	var startFrame = []byte{0x00, 0x00, 0x00, 0x00}

	packet.Write(startFrame)

	// data
	for _, led := range leds {
		var ledFrame = led.render(renderTime)

		packet.Write(ledFrame)
	}

	// stop
	var stopFrame = bytes.Repeat(stopByte, stopCount)

	packet.Write(stopFrame)

	if spiled.options.Debug {
		fmt.Println(hex.Dump(packet.Bytes()))
	}

	// send
	if _, err := spiled.spiBus.Write(packet.Bytes()); err != nil {
		return fmt.Errorf("SPIBUs.Write %v: %v", packet.Len(), err)
	} else {
		return nil
	}
}

func (spiled *SPILED) close() {
	defer spiled.waitGroup.Done()

	log.Printf("SPI-LED: Close...")

	if err := spiled.spiBus.Close(); err != nil {
		log.Printf("embd.SPIBus.Close: %v", err)
	}
}

func (spiled *SPILED) updateTally(tallyState tally.State) {
	log.Printf("SPI-LED: Update tally State:")

	leds := make([]LED, spiled.count)

	var found, errors int

	for i, led := range leds {
		if i == 0 {
			// skip status
			continue
		}

		id := tally.ID(i)

		if tally, exists := tallyState.Tally[id]; !exists {
			// missing tally state for pin
		} else {
			found++

			if tally.Status.Program && tally.Status.Preview {
				led = spiled.options.TallyBoth
			} else if tally.Status.Preview {
				led = spiled.options.TallyPreview
			} else if tally.Status.Program {
				led = spiled.options.TallyProgram
			} else {
				led = spiled.options.TallyIdle
			}

			led.Intensity = spiled.options.Intensity

			if tally.Errors != nil {
				led.Strobe(1 * time.Second)
			}

			log.Printf("SPI-LED %v: id=%v status=%v errors=%v led=%v", i, id, tally.Status, len(tally.Errors), led)

		}

		leds[i] = led
	}

	errors = len(tallyState.Errors)

	// status LED
	var statusLED LED

	if found > 0 && errors > 0 {
		statusLED = spiled.options.StatusWarn
	} else if errors > 0 {
		statusLED = spiled.options.StatusError
	} else if found > 0 {
		statusLED = spiled.options.StatusOK
	} else {
		statusLED = spiled.options.StatusIdle
	}

	statusLED.Intensity = spiled.options.Intensity

	log.Printf("SPI-LED: found=%v errors=%v led=%v", found, errors, statusLED)

	leds[0] = statusLED

	// refresh
	spiled.leds = leds
}

func (spiled *SPILED) run() {
	defer spiled.close()

	refreshTimer := time.Tick(time.Duration(1.0 / spiled.options.Refresh * float64(time.Second)))

	for {
		select {
		case tallyState, ok := <-spiled.tallyChan:
			if ok {
				spiled.updateTally(tallyState)
			} else {
				return
			}
		case refreshTime := <-refreshTimer:
			if err := spiled.write(spiled.leds, refreshTime); err != nil {
				log.Printf("SPI-LED: Write error: %v", err)
			}
		}
	}

	log.Printf("SPI-LED: Done")
}

func (spiled *SPILED) RegisterTally(t *tally.Tally) {
	spiled.tallyChan = make(chan tally.State)
	spiled.waitGroup.Add(1)

	go spiled.run()

	t.Register(spiled.tallyChan)
}

// Close and Wait..
func (spiled *SPILED) Close() {
	log.Printf("SPI-LED: Close..")

	if spiled.tallyChan != nil {
		close(spiled.tallyChan)
	}

	spiled.waitGroup.Wait()
}
<|endoftext|>


<commit_before>package client

import (
	"fmt"
	"strconv"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/docker/go-units"

	gflag "github.com/jessevdk/go-flags"
)

func (cli *HyperClient) HyperCmdImages(args ...string) error {
	var opts struct {
		All   bool `short:"a" long:"all" default:"false" description:"Show all images (by default filter out the intermediate image layers)"`
		Quiet bool `short:"q" long:"quiet" default:"false" description:"Only show numeric IDs"`
	}
	var parser = gflag.NewParser(&opts, gflag.Default)

	parser.Usage = "images [OPTIONS] [REPOSITORY]\n\nList images"
	args, err := parser.ParseArgs(args)
	if err != nil {
		if !strings.Contains(err.Error(), "Usage") {
			return err
		} else {
			return nil
		}
	}

	remoteInfo, err := cli.client.GetImages(opts.All, opts.Quiet)

	var (
		imagesList = []string{}
	)
	imagesList = remoteInfo.GetList("imagesList")

	w := tabwriter.NewWriter(cli.out, 20, 1, 3, ' ', 0)
	if opts.Quiet == false {
		fmt.Fprintln(w, "REPOSITORY\tTAG\tIMAGE ID\tCREATED\tVIRTUAL SIZE")
		for _, item := range imagesList {
			fields := strings.Split(item, ":")
			date, _ := strconv.ParseUint(fields[3], 0, 64)
			fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", fields[0], fields[1], fields[2][:12], time.Unix(int64(date), 0).Format("2006-01-02 15:04:05"), getImageSizeString(fields[4]))
		}
	} else {
		for _, item := range imagesList {
			fields := strings.Split(item, ":")
			fmt.Fprintf(w, "%s\n", fields[2][:12])
		}
	}
	w.Flush()

	return nil
}

func getImageSizeString(size string) string {
	s, err := strconv.Atoi(size)
	if err != nil {
		return ""
	}

	rtn := float64(s)
	return units.HumanSize(rtn)
}
<commit_msg>add err condition avoid panic<commit_after>package client

import (
	"fmt"
	"strconv"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/docker/go-units"

	gflag "github.com/jessevdk/go-flags"
)

func (cli *HyperClient) HyperCmdImages(args ...string) error {
	var opts struct {
		All   bool `short:"a" long:"all" default:"false" description:"Show all images (by default filter out the intermediate image layers)"`
		Quiet bool `short:"q" long:"quiet" default:"false" description:"Only show numeric IDs"`
	}
	var parser = gflag.NewParser(&opts, gflag.Default)

	parser.Usage = "images [OPTIONS] [REPOSITORY]\n\nList images"
	args, err := parser.ParseArgs(args)
	if err != nil {
		if !strings.Contains(err.Error(), "Usage") {
			return err
		} else {
			return nil
		}
	}

	remoteInfo, err := cli.client.GetImages(opts.All, opts.Quiet)
	if err != nil {
		return err
	}

	var (
		imagesList = []string{}
	)
	imagesList = remoteInfo.GetList("imagesList")

	w := tabwriter.NewWriter(cli.out, 20, 1, 3, ' ', 0)
	if opts.Quiet == false {
		fmt.Fprintln(w, "REPOSITORY\tTAG\tIMAGE ID\tCREATED\tVIRTUAL SIZE")
		for _, item := range imagesList {
			fields := strings.Split(item, ":")
			date, _ := strconv.ParseUint(fields[3], 0, 64)
			fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", fields[0], fields[1], fields[2][:12], time.Unix(int64(date), 0).Format("2006-01-02 15:04:05"), getImageSizeString(fields[4]))
		}
	} else {
		for _, item := range imagesList {
			fields := strings.Split(item, ":")
			fmt.Fprintf(w, "%s\n", fields[2][:12])
		}
	}
	w.Flush()

	return nil
}

func getImageSizeString(size string) string {
	s, err := strconv.Atoi(size)
	if err != nil {
		return ""
	}

	rtn := float64(s)
	return units.HumanSize(rtn)
}
<|endoftext|>


<commit_before>package regression

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"go.skia.org/infra/go/testutils/unittest"
	"go.skia.org/infra/perf/go/clustering2"
	"go.skia.org/infra/perf/go/dataframe"
	"go.skia.org/infra/perf/go/stepfit"
)

func TestRegressions(t *testing.T) {
	unittest.SmallTest(t)
	r := New()
	assert.True(t, r.Triaged(), "With no clusters, it should have Triaged() == true.")

	df := &dataframe.FrameResponse{}
	cl := &clustering2.ClusterSummary{}
	r.SetLow("source_type=skp", df, cl)
	assert.False(t, r.Triaged(), "Should not be Triaged.")

	// Triage the low cluster.
	err := r.TriageLow("source_type=skp", TriageStatus{
		Status:  POSITIVE,
		Message: "SKP Update",
	})
	assert.NoError(t, err)
	assert.True(t, r.Triaged())

	// Trying to triage a high cluster that doesn't exists.
	err = r.TriageHigh("source_type=skp", TriageStatus{
		Status: NEGATIVE,
	})
	assert.Equal(t, err, ErrNoClusterFound)
	assert.True(t, r.Triaged())

	// Set a high cluster.
	r.SetHigh("source_type=skp", df, cl)
	assert.False(t, r.Triaged())

	// And triage the high cluster.
	err = r.TriageHigh("source_type=skp", TriageStatus{
		Status:  NEGATIVE,
		Message: "See bug #foo.",
	})
	assert.NoError(t, err)
	assert.True(t, r.Triaged())

	// Trying to triage an unknown query.
	err = r.TriageHigh("uknownquery", TriageStatus{
		Status: NEGATIVE,
	})
	assert.Equal(t, err, ErrNoClusterFound)

	// Try serializing to JSON.
	b, err := r.JSON()
	assert.NoError(t, err)
	assert.Equal(t, "{\"by_query\":{\"source_type=skp\":{\"low\":{\"centroid\":null,\"shortcut\":\"\",\"param_summaries\":null,\"step_fit\":null,\"step_point\":null,\"num\":0},\"high\":{\"centroid\":null,\"shortcut\":\"\",\"param_summaries\":null,\"step_fit\":null,\"step_point\":null,\"num\":0},\"frame\":{\"dataframe\":null,\"ticks\":null,\"skps\":null,\"msg\":\"\"},\"low_status\":{\"status\":\"positive\",\"message\":\"SKP Update\"},\"high_status\":{\"status\":\"negative\",\"message\":\"See bug #foo.\"}}}}", string(b))
}

func TestMerge(t *testing.T) {
	unittest.SmallTest(t)

	r := newRegression()

	rhs := newRegression()
	df := &dataframe.FrameResponse{}
	cl := &clustering2.ClusterSummary{
		StepFit: &stepfit.StepFit{
			Regression: 1,
		},
	}
	rhs.Low = cl
	rhs.Frame = df

	r = r.Merge(rhs)
	assert.Equal(t, r.Low, cl)
	assert.Equal(t, r.Frame, df)

	r = r.Merge(rhs)
	assert.Equal(t, r.Low, cl)
	assert.Equal(t, r.Frame, df)

	clbetter := &clustering2.ClusterSummary{
		StepFit: &stepfit.StepFit{
			Regression: 2,
		},
	}
	dfbetter := &dataframe.FrameResponse{}
	betterlow := newRegression()
	betterlow.Low = clbetter
	betterlow.Frame = dfbetter

	r = r.Merge(betterlow)
	assert.Equal(t, r.Low, clbetter)
	assert.Equal(t, r.Frame, dfbetter)

	r = r.Merge(betterlow)
	assert.Equal(t, r.Low, clbetter)
	assert.Equal(t, r.Frame, dfbetter)

	// Now the same for High.
	rhs = newRegression()
	df = &dataframe.FrameResponse{}
	cl = &clustering2.ClusterSummary{
		StepFit: &stepfit.StepFit{
			Regression: -1,
		},
	}
	rhs.High = cl
	rhs.Frame = df

	r = r.Merge(rhs)
	assert.Equal(t, r.High, cl)
	assert.Equal(t, r.Frame, df)

	r = r.Merge(rhs)
	assert.Equal(t, r.High, cl)
	assert.Equal(t, r.Frame, df)

	clbetter = &clustering2.ClusterSummary{
		StepFit: &stepfit.StepFit{
			Regression: -2,
		},
	}
	dfbetter = &dataframe.FrameResponse{}
	betterhigh := newRegression()
	betterhigh.High = clbetter
	betterhigh.Frame = dfbetter

	r = r.Merge(betterhigh)
	assert.Equal(t, r.High, clbetter)
	assert.Equal(t, r.Frame, dfbetter)

	r = r.Merge(betterhigh)
	assert.Equal(t, r.High, clbetter)
	assert.Equal(t, r.Frame, dfbetter)
}
<commit_msg>[perf] Fix broken test.<commit_after>package regression

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"go.skia.org/infra/go/testutils/unittest"
	"go.skia.org/infra/perf/go/clustering2"
	"go.skia.org/infra/perf/go/dataframe"
	"go.skia.org/infra/perf/go/stepfit"
)

func TestRegressions(t *testing.T) {
	unittest.SmallTest(t)
	r := New()
	assert.True(t, r.Triaged(), "With no clusters, it should have Triaged() == true.")

	df := &dataframe.FrameResponse{}
	cl := &clustering2.ClusterSummary{}
	r.SetLow("source_type=skp", df, cl)
	assert.False(t, r.Triaged(), "Should not be Triaged.")

	// Triage the low cluster.
	err := r.TriageLow("source_type=skp", TriageStatus{
		Status:  POSITIVE,
		Message: "SKP Update",
	})
	assert.NoError(t, err)
	assert.True(t, r.Triaged())

	// Trying to triage a high cluster that doesn't exists.
	err = r.TriageHigh("source_type=skp", TriageStatus{
		Status: NEGATIVE,
	})
	assert.Equal(t, err, ErrNoClusterFound)
	assert.True(t, r.Triaged())

	// Set a high cluster.
	r.SetHigh("source_type=skp", df, cl)
	assert.False(t, r.Triaged())

	// And triage the high cluster.
	err = r.TriageHigh("source_type=skp", TriageStatus{
		Status:  NEGATIVE,
		Message: "See bug #foo.",
	})
	assert.NoError(t, err)
	assert.True(t, r.Triaged())

	// Trying to triage an unknown query.
	err = r.TriageHigh("uknownquery", TriageStatus{
		Status: NEGATIVE,
	})
	assert.Equal(t, err, ErrNoClusterFound)

	// Try serializing to JSON.
	b, err := r.JSON()
	assert.NoError(t, err)
	assert.Equal(t, "{\"by_query\":{\"source_type=skp\":{\"low\":{\"centroid\":null,\"shortcut\":\"\",\"param_summaries2\":null,\"step_fit\":null,\"step_point\":null,\"num\":0},\"high\":{\"centroid\":null,\"shortcut\":\"\",\"param_summaries2\":null,\"step_fit\":null,\"step_point\":null,\"num\":0},\"frame\":{\"dataframe\":null,\"ticks\":null,\"skps\":null,\"msg\":\"\"},\"low_status\":{\"status\":\"positive\",\"message\":\"SKP Update\"},\"high_status\":{\"status\":\"negative\",\"message\":\"See bug #foo.\"}}}}", string(b))
}

func TestMerge(t *testing.T) {
	unittest.SmallTest(t)

	r := newRegression()

	rhs := newRegression()
	df := &dataframe.FrameResponse{}
	cl := &clustering2.ClusterSummary{
		StepFit: &stepfit.StepFit{
			Regression: 1,
		},
	}
	rhs.Low = cl
	rhs.Frame = df

	r = r.Merge(rhs)
	assert.Equal(t, r.Low, cl)
	assert.Equal(t, r.Frame, df)

	r = r.Merge(rhs)
	assert.Equal(t, r.Low, cl)
	assert.Equal(t, r.Frame, df)

	clbetter := &clustering2.ClusterSummary{
		StepFit: &stepfit.StepFit{
			Regression: 2,
		},
	}
	dfbetter := &dataframe.FrameResponse{}
	betterlow := newRegression()
	betterlow.Low = clbetter
	betterlow.Frame = dfbetter

	r = r.Merge(betterlow)
	assert.Equal(t, r.Low, clbetter)
	assert.Equal(t, r.Frame, dfbetter)

	r = r.Merge(betterlow)
	assert.Equal(t, r.Low, clbetter)
	assert.Equal(t, r.Frame, dfbetter)

	// Now the same for High.
	rhs = newRegression()
	df = &dataframe.FrameResponse{}
	cl = &clustering2.ClusterSummary{
		StepFit: &stepfit.StepFit{
			Regression: -1,
		},
	}
	rhs.High = cl
	rhs.Frame = df

	r = r.Merge(rhs)
	assert.Equal(t, r.High, cl)
	assert.Equal(t, r.Frame, df)

	r = r.Merge(rhs)
	assert.Equal(t, r.High, cl)
	assert.Equal(t, r.Frame, df)

	clbetter = &clustering2.ClusterSummary{
		StepFit: &stepfit.StepFit{
			Regression: -2,
		},
	}
	dfbetter = &dataframe.FrameResponse{}
	betterhigh := newRegression()
	betterhigh.High = clbetter
	betterhigh.Frame = dfbetter

	r = r.Merge(betterhigh)
	assert.Equal(t, r.High, clbetter)
	assert.Equal(t, r.Frame, dfbetter)

	r = r.Merge(betterhigh)
	assert.Equal(t, r.High, clbetter)
	assert.Equal(t, r.Frame, dfbetter)
}
<|endoftext|>


<commit_before>package ground

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"errors"
	"fmt"
	"math"
	"math/big"
	"net"
	"time"

	certutil "k8s.io/client-go/util/cert"
)

const (
	duration365d = time.Hour * 24 * 365
)

type Certificates struct {
	Etcd struct {
		Clients struct {
			CA        *Bundle
			Apiserver *Bundle
		}
	}

	Kubernetes struct {
		Clients struct {
			CA                *Bundle
			ControllerManager *Bundle
			Scheduler         *Bundle
			Proxy             *Bundle
			Kubelet           *Bundle
			ClusterAdmin      *Bundle
		}
		Nodes struct {
			CA      *Bundle
			Generic *Bundle
		}
	}

	TLS struct {
		CA        *Bundle
		ApiServer *Bundle
		Etcd      *Bundle
	}
}

type Bundle struct {
	Certificate *x509.Certificate
	PrivateKey  *rsa.PrivateKey
}

type Config struct {
	CommonName         string
	Organization       []string
	OrganizationalUnit []string
	AltNames           AltNames
	Usages             []x509.ExtKeyUsage
}

type AltNames struct {
	DNSNames []string
	IPs      []net.IP
}

func (c Certificates) all() []*Bundle {
	return []*Bundle{
		c.Etcd.Clients.Apiserver,
		c.Etcd.Clients.CA,
		c.Kubernetes.Clients.CA,
		c.Kubernetes.Clients.ControllerManager,
		c.Kubernetes.Clients.Scheduler,
		c.Kubernetes.Clients.Proxy,
		c.Kubernetes.Clients.Kubelet,
		c.Kubernetes.Clients.ClusterAdmin,
		c.Kubernetes.Nodes.CA,
		c.Kubernetes.Nodes.Generic,
		c.TLS.CA,
		c.TLS.ApiServer,
		c.TLS.Etcd,
	}
}

func newCertificates(satellite string) (*Certificates, error) {
	certs := &Certificates{}

	if ca, err := newCA(satellite, "Etcd Clients"); err != nil {
		return certs, err
	} else {
		certs.Etcd.Clients.CA = ca
	}

	if ca, err := newCA(satellite, "Kubernetes Clients"); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Clients.CA = ca
	}

	if ca, err := newCA(satellite, "Kubernetes Nodes"); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Nodes.CA = ca
	}

	if ca, err := newCA(satellite, "TLS"); err != nil {
		return certs, err
	} else {
		certs.TLS.CA = ca
	}

	if cert, err := newSignedBundle(Config{
		CommonName: "apiserver",
		Usages:     []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, certs.Etcd.Clients.CA); err != nil {
		return certs, err
	} else {
		certs.Etcd.Clients.Apiserver = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName:   "cluster-admin",
		Organization: []string{"system:masters"},
		Usages:       []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, certs.Kubernetes.Clients.CA); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Clients.ClusterAdmin = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName: "system:kube-controller-manager",
		Usages:     []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, certs.Kubernetes.Clients.CA); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Clients.ControllerManager = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName:   "kuelet",
		Organization: []string{"system:nodes"},
		Usages:       []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, certs.Kubernetes.Clients.CA); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Clients.Kubelet = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName: "system:kube-proxy",
		Usages:     []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, certs.Kubernetes.Clients.CA); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Clients.Proxy = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName: "system:kube-scheduler",
		Usages:     []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, certs.Kubernetes.Clients.CA); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Clients.Scheduler = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName:   "kubelet",
		Organization: []string{"system:nodes"},
		Usages:       []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, certs.Kubernetes.Nodes.CA); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Nodes.Generic = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName: "apiserver",
		Usages:     []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		AltNames: AltNames{
			DNSNames: []string{"kubernetes", "kubernetes.default", "apiserver", "TODO:external.dns.name"},
			IPs:      []net.IP{net.IPv4(127, 0, 0, 1)},
		},
	}, certs.TLS.CA); err != nil {
		return certs, err
	} else {
		certs.TLS.ApiServer = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName: "etcd",
		Usages:     []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		AltNames: AltNames{
			DNSNames: []string{"etcd"},
			IPs:      []net.IP{net.IPv4(127, 0, 0, 1)},
		},
	}, certs.TLS.CA); err != nil {
		return certs, err
	} else {
		certs.TLS.Etcd = cert
	}

	return certs, nil
}

func newCA(satellite, name string) (*Bundle, error) {
	key, err := certutil.NewPrivateKey()
	if err != nil {
		panic(err)
		return &Bundle{}, fmt.Errorf("unable to create private key [%v]", err)
	}

	config := Config{
		CommonName:         name,
		OrganizationalUnit: []string{"SAP Converged Cloud", "Kubernikus", satellite},
	}
	cert, err := NewSelfSignedCACert(config, key)

	if err != nil {
		panic(err)
		return &Bundle{}, fmt.Errorf("unable to create self-signed certificate [%v]", err)
	}

	return &Bundle{cert, key}, nil
}

func newSignedBundle(config Config, ca *Bundle) (*Bundle, error) {
	key, err := certutil.NewPrivateKey()
	if err != nil {
		return &Bundle{}, fmt.Errorf("unable to create private key [%v]", err)
	}

	config.OrganizationalUnit = ca.Certificate.Subject.OrganizationalUnit

	cert, err := NewSignedCert(config, key, ca.Certificate, ca.PrivateKey)

	if err != nil {
		return &Bundle{}, fmt.Errorf("unable to create self-signed certificate [%v]", err)
	}

	return &Bundle{cert, key}, nil
}

func NewSelfSignedCACert(cfg Config, key *rsa.PrivateKey) (*x509.Certificate, error) {
	now := time.Now()
	tmpl := x509.Certificate{
		SerialNumber: new(big.Int).SetInt64(0),
		Subject: pkix.Name{
			CommonName:         cfg.CommonName,
			Organization:       cfg.Organization,
			OrganizationalUnit: cfg.OrganizationalUnit,
		},
		NotBefore:             now.UTC(),
		NotAfter:              now.Add(duration365d * 10).UTC(),
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		BasicConstraintsValid: true,
		IsCA: true,
	}

	certDERBytes, err := x509.CreateCertificate(rand.Reader, &tmpl, &tmpl, key.Public(), key)
	if err != nil {
		return nil, err
	}
	return x509.ParseCertificate(certDERBytes)
}

func NewSignedCert(cfg Config, key *rsa.PrivateKey, caCert *x509.Certificate, caKey *rsa.PrivateKey) (*x509.Certificate, error) {
	serial, err := rand.Int(rand.Reader, new(big.Int).SetInt64(math.MaxInt64))
	if err != nil {
		return nil, err
	}
	if len(cfg.CommonName) == 0 {
		return nil, errors.New("must specify a CommonName")
	}
	if len(cfg.Usages) == 0 {
		return nil, errors.New("must specify at least one ExtKeyUsage")
	}

	certTmpl := x509.Certificate{
		Subject: pkix.Name{
			CommonName:         cfg.CommonName,
			Organization:       cfg.Organization,
			OrganizationalUnit: cfg.OrganizationalUnit,
		},
		DNSNames:     cfg.AltNames.DNSNames,
		IPAddresses:  cfg.AltNames.IPs,
		SerialNumber: serial,
		NotBefore:    caCert.NotBefore,
		NotAfter:     time.Now().Add(duration365d * 10).UTC(),
		KeyUsage:     x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  cfg.Usages,
	}
	certDERBytes, err := x509.CreateCertificate(rand.Reader, &certTmpl, caCert, key.Public(), caKey)
	if err != nil {
		return nil, err
	}
	return x509.ParseCertificate(certDERBytes)
}
<commit_msg>add etcd peer certificates<commit_after>package ground

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"errors"
	"fmt"
	"math"
	"math/big"
	"net"
	"time"

	certutil "k8s.io/client-go/util/cert"
)

const (
	duration365d = time.Hour * 24 * 365
)

type Certificates struct {
	Etcd struct {
		Clients struct {
			CA        *Bundle
			Apiserver *Bundle
		}
		Peers struct {
			CA        *Bundle
			Universal *Bundle
		}
	}

	Kubernetes struct {
		Clients struct {
			CA                *Bundle
			ControllerManager *Bundle
			Scheduler         *Bundle
			Proxy             *Bundle
			Kubelet           *Bundle
			ClusterAdmin      *Bundle
		}
		Nodes struct {
			CA        *Bundle
			Universal *Bundle
		}
	}

	TLS struct {
		CA        *Bundle
		ApiServer *Bundle
		Etcd      *Bundle
	}
}

type Bundle struct {
	Certificate *x509.Certificate
	PrivateKey  *rsa.PrivateKey
}

type Config struct {
	CommonName         string
	Organization       []string
	OrganizationalUnit []string
	AltNames           AltNames
	Usages             []x509.ExtKeyUsage
}

type AltNames struct {
	DNSNames []string
	IPs      []net.IP
}

func (c Certificates) all() []*Bundle {
	return []*Bundle{
		c.Etcd.Clients.Apiserver,
		c.Etcd.Clients.CA,
		c.Etcd.Peers.Universal,
		c.Etcd.Peers.CA,
		c.Kubernetes.Clients.CA,
		c.Kubernetes.Clients.ControllerManager,
		c.Kubernetes.Clients.Scheduler,
		c.Kubernetes.Clients.Proxy,
		c.Kubernetes.Clients.Kubelet,
		c.Kubernetes.Clients.ClusterAdmin,
		c.Kubernetes.Nodes.CA,
		c.Kubernetes.Nodes.Universal,
		c.TLS.CA,
		c.TLS.ApiServer,
		c.TLS.Etcd,
	}
}

func newCertificates(satellite string) (*Certificates, error) {
	certs := &Certificates{}

	if ca, err := newCA(satellite, "Etcd Clients"); err != nil {
		return certs, err
	} else {
		certs.Etcd.Clients.CA = ca
	}

	if ca, err := newCA(satellite, "Etcd Peers"); err != nil {
		return certs, err
	} else {
		certs.Etcd.Peers.CA = ca
	}

	if ca, err := newCA(satellite, "Kubernetes Clients"); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Clients.CA = ca
	}

	if ca, err := newCA(satellite, "Kubernetes Nodes"); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Nodes.CA = ca
	}

	if ca, err := newCA(satellite, "TLS"); err != nil {
		return certs, err
	} else {
		certs.TLS.CA = ca
	}

	if cert, err := newSignedBundle(Config{
		CommonName: "apiserver",
		Usages:     []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, certs.Etcd.Clients.CA); err != nil {
		return certs, err
	} else {
		certs.Etcd.Clients.Apiserver = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName: "universal",
		Usages:     []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
	}, certs.Etcd.Peers.CA); err != nil {
		return certs, err
	} else {
		certs.Etcd.Peers.Universal = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName:   "cluster-admin",
		Organization: []string{"system:masters"},
		Usages:       []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, certs.Kubernetes.Clients.CA); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Clients.ClusterAdmin = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName: "system:kube-controller-manager",
		Usages:     []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, certs.Kubernetes.Clients.CA); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Clients.ControllerManager = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName:   "kubelet",
		Organization: []string{"system:nodes"},
		Usages:       []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, certs.Kubernetes.Clients.CA); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Clients.Kubelet = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName: "system:kube-proxy",
		Usages:     []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, certs.Kubernetes.Clients.CA); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Clients.Proxy = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName: "system:kube-scheduler",
		Usages:     []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, certs.Kubernetes.Clients.CA); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Clients.Scheduler = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName:   "universal",
		Organization: []string{"system:nodes"},
		Usages:       []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}, certs.Kubernetes.Nodes.CA); err != nil {
		return certs, err
	} else {
		certs.Kubernetes.Nodes.Universal = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName: "apiserver",
		Usages:     []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		AltNames: AltNames{
			DNSNames: []string{"kubernetes", "kubernetes.default", "apiserver", "TODO:external.dns.name"},
			IPs:      []net.IP{net.IPv4(127, 0, 0, 1)},
		},
	}, certs.TLS.CA); err != nil {
		return certs, err
	} else {
		certs.TLS.ApiServer = cert
	}

	if cert, err := newSignedBundle(Config{
		CommonName: "etcd",
		Usages:     []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		AltNames: AltNames{
			DNSNames: []string{"etcd"},
			IPs:      []net.IP{net.IPv4(127, 0, 0, 1)},
		},
	}, certs.TLS.CA); err != nil {
		return certs, err
	} else {
		certs.TLS.Etcd = cert
	}

	return certs, nil
}

func newCA(satellite, name string) (*Bundle, error) {
	key, err := certutil.NewPrivateKey()
	if err != nil {
		panic(err)
		return &Bundle{}, fmt.Errorf("unable to create private key [%v]", err)
	}

	config := Config{
		CommonName:         name,
		OrganizationalUnit: []string{"SAP Converged Cloud", "Kubernikus", satellite},
	}
	cert, err := NewSelfSignedCACert(config, key)

	if err != nil {
		panic(err)
		return &Bundle{}, fmt.Errorf("unable to create self-signed certificate [%v]", err)
	}

	return &Bundle{cert, key}, nil
}

func newSignedBundle(config Config, ca *Bundle) (*Bundle, error) {
	key, err := certutil.NewPrivateKey()
	if err != nil {
		return &Bundle{}, fmt.Errorf("unable to create private key [%v]", err)
	}

	config.OrganizationalUnit = ca.Certificate.Subject.OrganizationalUnit

	cert, err := NewSignedCert(config, key, ca.Certificate, ca.PrivateKey)

	if err != nil {
		return &Bundle{}, fmt.Errorf("unable to create self-signed certificate [%v]", err)
	}

	return &Bundle{cert, key}, nil
}

func NewSelfSignedCACert(cfg Config, key *rsa.PrivateKey) (*x509.Certificate, error) {
	now := time.Now()
	tmpl := x509.Certificate{
		SerialNumber: new(big.Int).SetInt64(0),
		Subject: pkix.Name{
			CommonName:         cfg.CommonName,
			Organization:       cfg.Organization,
			OrganizationalUnit: cfg.OrganizationalUnit,
		},
		NotBefore:             now.UTC(),
		NotAfter:              now.Add(duration365d * 10).UTC(),
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		BasicConstraintsValid: true,
		IsCA: true,
	}

	certDERBytes, err := x509.CreateCertificate(rand.Reader, &tmpl, &tmpl, key.Public(), key)
	if err != nil {
		return nil, err
	}
	return x509.ParseCertificate(certDERBytes)
}

func NewSignedCert(cfg Config, key *rsa.PrivateKey, caCert *x509.Certificate, caKey *rsa.PrivateKey) (*x509.Certificate, error) {
	serial, err := rand.Int(rand.Reader, new(big.Int).SetInt64(math.MaxInt64))
	if err != nil {
		return nil, err
	}
	if len(cfg.CommonName) == 0 {
		return nil, errors.New("must specify a CommonName")
	}
	if len(cfg.Usages) == 0 {
		return nil, errors.New("must specify at least one ExtKeyUsage")
	}

	certTmpl := x509.Certificate{
		Subject: pkix.Name{
			CommonName:         cfg.CommonName,
			Organization:       cfg.Organization,
			OrganizationalUnit: cfg.OrganizationalUnit,
		},
		DNSNames:     cfg.AltNames.DNSNames,
		IPAddresses:  cfg.AltNames.IPs,
		SerialNumber: serial,
		NotBefore:    caCert.NotBefore,
		NotAfter:     time.Now().Add(duration365d * 10).UTC(),
		KeyUsage:     x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  cfg.Usages,
	}
	certDERBytes, err := x509.CreateCertificate(rand.Reader, &certTmpl, caCert, key.Public(), caKey)
	if err != nil {
		return nil, err
	}
	return x509.ParseCertificate(certDERBytes)
}
<|endoftext|>


<commit_before>// +build windows

package windows

import (
	"github.com/mackerelio/mackerel-agent/logging"
	"github.com/mackerelio/mackerel-agent/metrics"
	"github.com/mackerelio/mackerel-agent/util"
)

// CPUUsageGenerator XXX
type CPUUsageGenerator struct {
}

var cpuUsageLogger = logging.GetLogger("metrics.cpuusage")

// NewCPUUsageGenerator XXX
func NewCPUUsageGenerator() (*CPUUsageGenerator, error) {
	return &CPUUsageGenerator{}, nil
}

// Generate XXX
func (g *CPUUsageGenerator) Generate() (metrics.Values, error) {
	cpuusage := make(map[string]float64, 1)

	cpuusageValue, err := util.GetWmicToFloat("cpu", "loadpercentage")
	if err != nil {
		cpuusageValue = 0
	}
	cpuusage["cpu.user.percentage"] = cpuusageValue
	cpuusage["cpu.idle.percentage"] = 100 - cpuusageValue
	cpuUsageLogger.Debugf("cpuusage : %s", cpuusage)
	return metrics.Values(cpuusage), nil
}
<commit_msg>omit parameter<commit_after>// +build windows

package windows

import (
	"github.com/mackerelio/mackerel-agent/logging"
	"github.com/mackerelio/mackerel-agent/metrics"
	"github.com/mackerelio/mackerel-agent/util"
)

// CPUUsageGenerator XXX
type CPUUsageGenerator struct {
}

var cpuUsageLogger = logging.GetLogger("metrics.cpuusage")

// NewCPUUsageGenerator XXX
func NewCPUUsageGenerator() (*CPUUsageGenerator, error) {
	return &CPUUsageGenerator{}, nil
}

// Generate XXX
func (g *CPUUsageGenerator) Generate() (metrics.Values, error) {
	cpuusage := make(map[string]float64)

	cpuusageValue, err := util.GetWmicToFloat("cpu", "loadpercentage")
	if err != nil {
		cpuusageValue = 0
	}
	cpuusage["cpu.user.percentage"] = cpuusageValue
	cpuusage["cpu.idle.percentage"] = 100 - cpuusageValue
	cpuUsageLogger.Debugf("cpuusage : %s", cpuusage)
	return metrics.Values(cpuusage), nil
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package flowcontrol

import (
	"math"
	"sync"
	"testing"
	"time"
)

func TestMultithreadedThrottling(t *testing.T) {
	// Bucket with 100QPS and no burst
	r := NewTokenBucketRateLimiter(100, 1)

	// channel to collect 100 tokens
	taken := make(chan bool, 100)

	// Set up goroutines to hammer the throttler
	startCh := make(chan bool)
	endCh := make(chan bool)
	for i := 0; i < 10; i++ {
		go func() {
			// wait for the starting signal
			<-startCh
			for {
				// get a token
				r.Accept()
				select {
				// try to add it to the taken channel
				case taken <- true:
					continue
				// if taken is full, notify and return
				default:
					endCh <- true
					return
				}
			}
		}()
	}

	// record wall time
	startTime := time.Now()
	// take the initial capacity so all tokens are the result of refill
	r.Accept()
	// start the thundering herd
	close(startCh)
	// wait for the first signal that we collected 100 tokens
	<-endCh
	// record wall time
	endTime := time.Now()

	if duration := endTime.Sub(startTime); duration < time.Second {
		// We shouldn't be able to get 100 tokens out of the bucket in less than 1 second of wall clock time, no matter what
		t.Errorf("Expected it to take at least 1 second to get 100 tokens, took %v", duration)
	} else {
		t.Logf("Took %v to get 100 tokens", duration)
	}
}

func TestBasicThrottle(t *testing.T) {
	r := NewTokenBucketRateLimiter(1, 3)
	for i := 0; i < 3; i++ {
		if !r.TryAccept() {
			t.Error("unexpected false accept")
		}
	}
	if r.TryAccept() {
		t.Error("unexpected true accept")
	}
}

func TestIncrementThrottle(t *testing.T) {
	r := NewTokenBucketRateLimiter(1, 1)
	if !r.TryAccept() {
		t.Error("unexpected false accept")
	}
	if r.TryAccept() {
		t.Error("unexpected true accept")
	}

	// Allow to refill
	time.Sleep(2 * time.Second)

	if !r.TryAccept() {
		t.Error("unexpected false accept")
	}
}

func TestThrottle(t *testing.T) {
	r := NewTokenBucketRateLimiter(10, 5)

	// Should consume 5 tokens immediately, then
	// the remaining 11 should take at least 1 second (0.1s each)
	expectedFinish := time.Now().Add(time.Second * 1)
	for i := 0; i < 16; i++ {
		r.Accept()
	}
	if time.Now().Before(expectedFinish) {
		t.Error("rate limit was not respected, finished too early")
	}
}

func TestRateLimiterSaturation(t *testing.T) {
	const e = 0.000001
	tests := []struct {
		capacity int
		take     int

		expectedSaturation float64
	}{
		{1, 1, 1},
		{10, 3, 0.3},
	}
	for i, tt := range tests {
		rl := NewTokenBucketRateLimiter(1, tt.capacity)
		for i := 0; i < tt.take; i++ {
			rl.Accept()
		}
		if math.Abs(rl.Saturation()-tt.expectedSaturation) > e {
			t.Fatalf("#%d: Saturation rate difference isn't within tolerable range\n want=%f, get=%f",
				i, tt.expectedSaturation, rl.Saturation())
		}
	}
}

func TestAlwaysFake(t *testing.T) {
	rl := NewFakeAlwaysRateLimiter()
	if !rl.TryAccept() {
		t.Error("TryAccept in AlwaysFake should return true.")
	}
	// If this will block the test will timeout
	rl.Accept()
}

func TestNeverFake(t *testing.T) {
	rl := NewFakeNeverRateLimiter()
	if rl.TryAccept() {
		t.Error("TryAccept in NeverFake should return false.")
	}

	finished := false
	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		rl.Accept()
		finished = true
		wg.Done()
	}()

	// Wait some time to make sure it never finished.
	time.Sleep(time.Second)
	if finished {
		t.Error("Accept should block forever in NeverFake.")
	}

	rl.Stop()
	wg.Wait()
	if !finished {
		t.Error("Stop should make Accept unblock in NeverFake.")
	}
}
<commit_msg>tolerate clock change in throttle testing<commit_after>/*
Copyright 2014 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package flowcontrol

import (
	"math"
	"sync"
	"testing"
	"time"
)

func TestMultithreadedThrottling(t *testing.T) {
	// Bucket with 100QPS and no burst
	r := NewTokenBucketRateLimiter(100, 1)

	// channel to collect 100 tokens
	taken := make(chan bool, 100)

	// Set up goroutines to hammer the throttler
	startCh := make(chan bool)
	endCh := make(chan bool)
	for i := 0; i < 10; i++ {
		go func() {
			// wait for the starting signal
			<-startCh
			for {
				// get a token
				r.Accept()
				select {
				// try to add it to the taken channel
				case taken <- true:
					continue
				// if taken is full, notify and return
				default:
					endCh <- true
					return
				}
			}
		}()
	}

	// record wall time
	startTime := time.Now()
	// take the initial capacity so all tokens are the result of refill
	r.Accept()
	// start the thundering herd
	close(startCh)
	// wait for the first signal that we collected 100 tokens
	<-endCh
	// record wall time
	endTime := time.Now()

	// tolerate a 1% clock change because these things happen
	if duration := endTime.Sub(startTime); duration < (time.Second * 99 / 100) {
		// We shouldn't be able to get 100 tokens out of the bucket in less than 1 second of wall clock time, no matter what
		t.Errorf("Expected it to take at least 1 second to get 100 tokens, took %v", duration)
	} else {
		t.Logf("Took %v to get 100 tokens", duration)
	}
}

func TestBasicThrottle(t *testing.T) {
	r := NewTokenBucketRateLimiter(1, 3)
	for i := 0; i < 3; i++ {
		if !r.TryAccept() {
			t.Error("unexpected false accept")
		}
	}
	if r.TryAccept() {
		t.Error("unexpected true accept")
	}
}

func TestIncrementThrottle(t *testing.T) {
	r := NewTokenBucketRateLimiter(1, 1)
	if !r.TryAccept() {
		t.Error("unexpected false accept")
	}
	if r.TryAccept() {
		t.Error("unexpected true accept")
	}

	// Allow to refill
	time.Sleep(2 * time.Second)

	if !r.TryAccept() {
		t.Error("unexpected false accept")
	}
}

func TestThrottle(t *testing.T) {
	r := NewTokenBucketRateLimiter(10, 5)

	// Should consume 5 tokens immediately, then
	// the remaining 11 should take at least 1 second (0.1s each)
	expectedFinish := time.Now().Add(time.Second * 1)
	for i := 0; i < 16; i++ {
		r.Accept()
	}
	if time.Now().Before(expectedFinish) {
		t.Error("rate limit was not respected, finished too early")
	}
}

func TestRateLimiterSaturation(t *testing.T) {
	const e = 0.000001
	tests := []struct {
		capacity int
		take     int

		expectedSaturation float64
	}{
		{1, 1, 1},
		{10, 3, 0.3},
	}
	for i, tt := range tests {
		rl := NewTokenBucketRateLimiter(1, tt.capacity)
		for i := 0; i < tt.take; i++ {
			rl.Accept()
		}
		if math.Abs(rl.Saturation()-tt.expectedSaturation) > e {
			t.Fatalf("#%d: Saturation rate difference isn't within tolerable range\n want=%f, get=%f",
				i, tt.expectedSaturation, rl.Saturation())
		}
	}
}

func TestAlwaysFake(t *testing.T) {
	rl := NewFakeAlwaysRateLimiter()
	if !rl.TryAccept() {
		t.Error("TryAccept in AlwaysFake should return true.")
	}
	// If this will block the test will timeout
	rl.Accept()
}

func TestNeverFake(t *testing.T) {
	rl := NewFakeNeverRateLimiter()
	if rl.TryAccept() {
		t.Error("TryAccept in NeverFake should return false.")
	}

	finished := false
	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		rl.Accept()
		finished = true
		wg.Done()
	}()

	// Wait some time to make sure it never finished.
	time.Sleep(time.Second)
	if finished {
		t.Error("Accept should block forever in NeverFake.")
	}

	rl.Stop()
	wg.Wait()
	if !finished {
		t.Error("Stop should make Accept unblock in NeverFake.")
	}
}
<|endoftext|>


<commit_before>package schema

import (
	"bufio"
	"bytes"
	"camli/blobref"
	"crypto/sha1"
	"fmt"
	"io"
	"json"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

var NoCamliVersionError = os.NewError("No camliVersion key in map")
var UnimplementedError = os.NewError("Unimplemented")

type StatHasher interface {
	Lstat(fileName string) (*os.FileInfo, os.Error)
	Hash(fileName string) (*blobref.BlobRef, os.Error)
}

var DefaultStatHasher = &defaultStatHasher{}

type defaultStatHasher struct{}

func (d *defaultStatHasher) Lstat(fileName string) (*os.FileInfo, os.Error) {
	return os.Lstat(fileName)
}

func (d *defaultStatHasher) Hash(fileName string) (*blobref.BlobRef, os.Error) {
	s1 := sha1.New()
	file, err := os.Open(fileName, os.O_RDONLY, 0)
	if err != nil {
		return nil, err
	}
	defer file.Close()
	_, err = io.Copy(s1, file)
        if err != nil {
                return nil, err
        }
	return blobref.FromHash("sha1", s1), nil
}

type StaticSet struct {
	l  sync.Mutex
	refs []*blobref.BlobRef
}

func (ss *StaticSet) Add(ref *blobref.BlobRef) {
	ss.l.Lock()
	defer ss.l.Unlock()
	ss.refs = append(ss.refs, ref)
}

func newCamliMap(version int, ctype string) map[string]interface{} {
	m := make(map[string]interface{})
        m["camliVersion"] = version
        m["camliType"] = ctype
	return m
}

// Map returns a Camli map of camliType "static-set"
func (ss *StaticSet) Map() map[string]interface{} {
	m := newCamliMap(1, "static-set")
	ss.l.Lock()
	defer ss.l.Unlock()

	members := make([]string, 0, len(ss.refs))
	if ss.refs != nil {
		for _, ref := range ss.refs {
			members = append(members, ref.String())
		}
	}
	m["members"] = members
	return m
}

func MapToCamliJson(m map[string]interface{}) (string, os.Error) {
	version, hasVersion := m["camliVersion"]
	if !hasVersion {
		return "", NoCamliVersionError
	}
	m["camliVersion"] = 0, false
	jsonBytes, err := json.MarshalIndent(m, "", "  ")
	if err != nil {
		return "", err
	}
	m["camliVersion"] = version
	buf := new(bytes.Buffer)
	fmt.Fprintf(buf, "{\"camliVersion\": %v,\n", version)
	buf.Write(jsonBytes[2:])
	return string(buf.Bytes()), nil
}

func NewCommonFileMap(fileName string, fi *os.FileInfo) map[string]interface{} {
	m := newCamliMap(1, "" /* no type yet */)

	lastSlash := strings.LastIndex(fileName, "/")
	baseName := fileName[lastSlash+1:]
	if isValidUtf8(baseName) {
		m["fileName"] = baseName
	} else {
		m["fileNameBytes"] = []uint8(baseName)
	}

	// Common elements (from file-common.txt)
	m["unixPermission"] = fmt.Sprintf("0%o", fi.Permission())
	if fi.Uid != -1 {
		m["unixOwnerId"] = fi.Uid
		if user := getUserFromUid(fi.Uid); user != "" {
			m["unixOwner"] = user
		}
	}
	if fi.Gid != -1 {
		m["unixGroupId"] = fi.Gid
		if group := getGroupFromGid(fi.Gid); group != "" {
			m["unixGroup"] = group
		}
	}
	if mtime := fi.Mtime_ns; mtime != 0 {
		m["unixMtime"] = rfc3339FromNanos(mtime)
	}

	return m
}

type ContentPart struct {
	BlobRef   *blobref.BlobRef
	Size      int64
	Offset    int64
}

type InvalidContentPartsError struct {
	StatSize   int64
	SumOfParts int64
}

func (e *InvalidContentPartsError) String() string {
	return fmt.Sprintf("Invalid ContentPart slice in PopulateRegularFileMap; file stat size is %d but sum of parts was %d", e.StatSize, e.SumOfParts)
}

func PopulateRegularFileMap(m map[string]interface{}, fi *os.FileInfo, parts []ContentPart) os.Error {
	m["camliType"] = "file"
	m["size"] = fi.Size

	sumSize := int64(0)
	mparts := make([]map[string]interface{}, len(parts))
	for idx, part := range parts {
		mpart := make(map[string]interface{})
		mparts[idx] = mpart
		mpart["blobRef"] = part.BlobRef.String()
		mpart["size"] = part.Size
		sumSize += part.Size
		if part.Offset != 0 {
			mpart["offset"] = part.Offset
		}
	}
	if sumSize != fi.Size {
		return &InvalidContentPartsError{fi.Size, sumSize}
	}
	m["contentParts"] = mparts
	return nil
}

func PopulateSymlinkMap(m map[string]interface{}, fileName string) os.Error {
	m["camliType"] = "symlink"
	target, err := os.Readlink(fileName)
	if err != nil {
		return err
	}
	if isValidUtf8(target) {
		m["symlinkTarget"] = target
	} else {
		m["symlinkTargetBytes"] = []uint8(target)
	}
	return nil
}

func PopulateDirectoryMap(m map[string]interface{}, staticSetRef *blobref.BlobRef) {
	m["camliType"] = "directory"
	m["entries"] = staticSetRef.String()
}

func rfc3339FromNanos(epochnanos int64) string {
	nanos := epochnanos % 1e9
	esec := epochnanos / 1e9
	t := time.SecondsToUTC(esec)
	timeStr := t.Format(time.RFC3339)
	if nanos == 0 {
		return timeStr
	}
	nanoStr := fmt.Sprintf("%09d", nanos)
	nanoStr = strings.TrimRight(nanoStr, "0")
	return timeStr[:len(timeStr)-1] + "." + nanoStr + "Z"
}

func populateMap(m map[int]string, file string) {
	f, err := os.Open(file, os.O_RDONLY, 0)
	if err != nil {
		return
	}
	bufr := bufio.NewReader(f)
	for {
		line, err := bufr.ReadString('\n')
		if err != nil {
			return
		}
		parts := strings.Split(line, ":", 4)
		if len(parts) >= 3 {
			idstr := parts[2]
			id, err := strconv.Atoi(idstr)
			if err == nil {
				m[id] = parts[0]
			}
		}
	}
}

var uidToUsernameMap map[int]string
var getUserFromUidOnce sync.Once
func getUserFromUid(uid int) string {
	getUserFromUidOnce.Do(func() {
		uidToUsernameMap = make(map[int]string)
		populateMap(uidToUsernameMap, "/etc/passwd")
	})
	return uidToUsernameMap[uid]
}

var gidToUsernameMap map[int]string
var getGroupFromGidOnce sync.Once
func getGroupFromGid(uid int) string {
	getGroupFromGidOnce.Do(func() {
		gidToUsernameMap = make(map[int]string)
		populateMap(gidToUsernameMap, "/etc/group")
	})
	return gidToUsernameMap[uid]
}

func isValidUtf8(s string) bool {
	for _, rune := range []int(s) {
		if rune == 0xfffd {
			return false
		}
	}
	return true
}
<commit_msg>Only include the ctime when it differs from mtime<commit_after>package schema

import (
	"bufio"
	"bytes"
	"camli/blobref"
	"crypto/sha1"
	"fmt"
	"io"
	"json"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

var NoCamliVersionError = os.NewError("No camliVersion key in map")
var UnimplementedError = os.NewError("Unimplemented")

type StatHasher interface {
	Lstat(fileName string) (*os.FileInfo, os.Error)
	Hash(fileName string) (*blobref.BlobRef, os.Error)
}

var DefaultStatHasher = &defaultStatHasher{}

type defaultStatHasher struct{}

func (d *defaultStatHasher) Lstat(fileName string) (*os.FileInfo, os.Error) {
	return os.Lstat(fileName)
}

func (d *defaultStatHasher) Hash(fileName string) (*blobref.BlobRef, os.Error) {
	s1 := sha1.New()
	file, err := os.Open(fileName, os.O_RDONLY, 0)
	if err != nil {
		return nil, err
	}
	defer file.Close()
	_, err = io.Copy(s1, file)
        if err != nil {
                return nil, err
        }
	return blobref.FromHash("sha1", s1), nil
}

type StaticSet struct {
	l  sync.Mutex
	refs []*blobref.BlobRef
}

func (ss *StaticSet) Add(ref *blobref.BlobRef) {
	ss.l.Lock()
	defer ss.l.Unlock()
	ss.refs = append(ss.refs, ref)
}

func newCamliMap(version int, ctype string) map[string]interface{} {
	m := make(map[string]interface{})
        m["camliVersion"] = version
        m["camliType"] = ctype
	return m
}

// Map returns a Camli map of camliType "static-set"
func (ss *StaticSet) Map() map[string]interface{} {
	m := newCamliMap(1, "static-set")
	ss.l.Lock()
	defer ss.l.Unlock()

	members := make([]string, 0, len(ss.refs))
	if ss.refs != nil {
		for _, ref := range ss.refs {
			members = append(members, ref.String())
		}
	}
	m["members"] = members
	return m
}

func MapToCamliJson(m map[string]interface{}) (string, os.Error) {
	version, hasVersion := m["camliVersion"]
	if !hasVersion {
		return "", NoCamliVersionError
	}
	m["camliVersion"] = 0, false
	jsonBytes, err := json.MarshalIndent(m, "", "  ")
	if err != nil {
		return "", err
	}
	m["camliVersion"] = version
	buf := new(bytes.Buffer)
	fmt.Fprintf(buf, "{\"camliVersion\": %v,\n", version)
	buf.Write(jsonBytes[2:])
	return string(buf.Bytes()), nil
}

func NewCommonFileMap(fileName string, fi *os.FileInfo) map[string]interface{} {
	m := newCamliMap(1, "" /* no type yet */)

	lastSlash := strings.LastIndex(fileName, "/")
	baseName := fileName[lastSlash+1:]
	if isValidUtf8(baseName) {
		m["fileName"] = baseName
	} else {
		m["fileNameBytes"] = []uint8(baseName)
	}

	// Common elements (from file-common.txt)
	m["unixPermission"] = fmt.Sprintf("0%o", fi.Permission())
	if fi.Uid != -1 {
		m["unixOwnerId"] = fi.Uid
		if user := getUserFromUid(fi.Uid); user != "" {
			m["unixOwner"] = user
		}
	}
	if fi.Gid != -1 {
		m["unixGroupId"] = fi.Gid
		if group := getGroupFromGid(fi.Gid); group != "" {
			m["unixGroup"] = group
		}
	}
	if mtime := fi.Mtime_ns; mtime != 0 {
		m["unixMtime"] = rfc3339FromNanos(mtime)
	}
	// Include the ctime too, if it differs.
	if ctime := fi.Ctime_ns; ctime != 0 && fi.Mtime_ns != fi.Ctime_ns {
		m["unixCtime"] = rfc3339FromNanos(ctime)
	}

	return m
}

type ContentPart struct {
	BlobRef   *blobref.BlobRef
	Size      int64
	Offset    int64
}

type InvalidContentPartsError struct {
	StatSize   int64
	SumOfParts int64
}

func (e *InvalidContentPartsError) String() string {
	return fmt.Sprintf("Invalid ContentPart slice in PopulateRegularFileMap; file stat size is %d but sum of parts was %d", e.StatSize, e.SumOfParts)
}

func PopulateRegularFileMap(m map[string]interface{}, fi *os.FileInfo, parts []ContentPart) os.Error {
	m["camliType"] = "file"
	m["size"] = fi.Size

	sumSize := int64(0)
	mparts := make([]map[string]interface{}, len(parts))
	for idx, part := range parts {
		mpart := make(map[string]interface{})
		mparts[idx] = mpart
		mpart["blobRef"] = part.BlobRef.String()
		mpart["size"] = part.Size
		sumSize += part.Size
		if part.Offset != 0 {
			mpart["offset"] = part.Offset
		}
	}
	if sumSize != fi.Size {
		return &InvalidContentPartsError{fi.Size, sumSize}
	}
	m["contentParts"] = mparts
	return nil
}

func PopulateSymlinkMap(m map[string]interface{}, fileName string) os.Error {
	m["camliType"] = "symlink"
	target, err := os.Readlink(fileName)
	if err != nil {
		return err
	}
	if isValidUtf8(target) {
		m["symlinkTarget"] = target
	} else {
		m["symlinkTargetBytes"] = []uint8(target)
	}
	return nil
}

func PopulateDirectoryMap(m map[string]interface{}, staticSetRef *blobref.BlobRef) {
	m["camliType"] = "directory"
	m["entries"] = staticSetRef.String()
}

func rfc3339FromNanos(epochnanos int64) string {
	nanos := epochnanos % 1e9
	esec := epochnanos / 1e9
	t := time.SecondsToUTC(esec)
	timeStr := t.Format(time.RFC3339)
	if nanos == 0 {
		return timeStr
	}
	nanoStr := fmt.Sprintf("%09d", nanos)
	nanoStr = strings.TrimRight(nanoStr, "0")
	return timeStr[:len(timeStr)-1] + "." + nanoStr + "Z"
}

func populateMap(m map[int]string, file string) {
	f, err := os.Open(file, os.O_RDONLY, 0)
	if err != nil {
		return
	}
	bufr := bufio.NewReader(f)
	for {
		line, err := bufr.ReadString('\n')
		if err != nil {
			return
		}
		parts := strings.Split(line, ":", 4)
		if len(parts) >= 3 {
			idstr := parts[2]
			id, err := strconv.Atoi(idstr)
			if err == nil {
				m[id] = parts[0]
			}
		}
	}
}

var uidToUsernameMap map[int]string
var getUserFromUidOnce sync.Once
func getUserFromUid(uid int) string {
	getUserFromUidOnce.Do(func() {
		uidToUsernameMap = make(map[int]string)
		populateMap(uidToUsernameMap, "/etc/passwd")
	})
	return uidToUsernameMap[uid]
}

var gidToUsernameMap map[int]string
var getGroupFromGidOnce sync.Once
func getGroupFromGid(uid int) string {
	getGroupFromGidOnce.Do(func() {
		gidToUsernameMap = make(map[int]string)
		populateMap(gidToUsernameMap, "/etc/group")
	})
	return gidToUsernameMap[uid]
}

func isValidUtf8(s string) bool {
	for _, rune := range []int(s) {
		if rune == 0xfffd {
			return false
		}
	}
	return true
}
<|endoftext|>