<commit_before>package grafana_test

import (
	"encoding/json"
	"fmt"

	"github.com/grafov/autograf/grafana"
)

func ExampleNewBoard() {
	board := grafana.NewBoard("Sample dashboard title")
	board.ID = 1
	row1 := board.AddRow("Sample row title")
	row1.Add(grafana.NewGraph("Sample graph"))
	graphWithDs := grafana.NewGraph("Sample graph 2")
	target := grafana.Target{
		RefID:      "A",
		Datasource: "Sample Source 1",
		Expr:       "sample request 1"}
	graphWithDs.AddTarget(&target)
	row1.Add(graphWithDs)
	data, _ := json.MarshalIndent(board, "", "  ")
	fmt.Printf("%s", data)
	// Output:
	// 	{
	//   "id": 1,
	//   "title": "Sample dashboard title",
	//   "originalTitle": "",
	//   "tags": null,
	//   "style": "dark",
	//   "timezone": "browser",
	//   "editable": true,
	//   "hideControls": false,
	//   "sharedCrosshair": false,
	//   "rows": [
	//     {
	//       "title": "Sample row title",
	//       "showTitle": false,
	//       "collapse": false,
	//       "editable": true,
	//       "height": "250px",
	//       "panels": [
	//         {
	//           "id": 1,
	//           "title": "Sample graph",
	//           "span": 12,
	//           "renderer": "flot",
	//           "transparent": false,
	//           "type": "graph",
	//           "error": false,
	//           "isNew": true,
	//           "aliasColors": null,
	//           "bars": false,
	//           "fill": 0,
	//           "grid": {
	//             "leftLogBase": null,
	//             "leftMax": null,
	//             "leftMin": null,
	//             "rightLogBase": null,
	//             "rightMax": null,
	//             "rightMin": null,
	//             "threshold1": null,
	//             "threshold1Color": "",
	//             "threshold2": null,
	//             "threshold2Color": "",
	//             "thresholdLine": false
	//           },
	//           "legend": {
	//             "alignAsTable": false,
	//             "avg": false,
	//             "current": false,
	//             "hideEmpty": false,
	//             "hideZero": false,
	//             "max": false,
	//             "min": false,
	//             "rightSide": false,
	//             "show": false,
	//             "total": false,
	//             "values": false
	//           },
	//           "lines": false,
	//           "linewidth": 0,
	//           "nullPointMode": "connected",
	//           "percentage": false,
	//           "pointradius": 5,
	//           "points": false,
	//           "seriesOverrides": null,
	//           "stack": false,
	//           "steppedLine": false,
	//           "timeFrom": null,
	//           "timeShift": null,
	//           "tooltip": {
	//             "shared": false,
	//             "value_type": ""
	//           },
	//           "x-axis": true,
	//           "y-axis": true,
	//           "y_formats": null
	//         },
	//         {
	//           "id": 2,
	//           "title": "Sample graph 2",
	//           "span": 12,
	//           "renderer": "flot",
	//           "transparent": false,
	//           "type": "graph",
	//           "error": false,
	//           "isNew": true,
	//           "aliasColors": null,
	//           "bars": false,
	//           "fill": 0,
	//           "grid": {
	//             "leftLogBase": null,
	//             "leftMax": null,
	//             "leftMin": null,
	//             "rightLogBase": null,
	//             "rightMax": null,
	//             "rightMin": null,
	//             "threshold1": null,
	//             "threshold1Color": "",
	//             "threshold2": null,
	//             "threshold2Color": "",
	//             "thresholdLine": false
	//           },
	//           "legend": {
	//             "alignAsTable": false,
	//             "avg": false,
	//             "current": false,
	//             "hideEmpty": false,
	//             "hideZero": false,
	//             "max": false,
	//             "min": false,
	//             "rightSide": false,
	//             "show": false,
	//             "total": false,
	//             "values": false
	//           },
	//           "lines": false,
	//           "linewidth": 0,
	//           "nullPointMode": "connected",
	//           "percentage": false,
	//           "pointradius": 5,
	//           "points": false,
	//           "seriesOverrides": null,
	//           "stack": false,
	//           "steppedLine": false,
	//           "targets": [
	//             {
	//               "refId": "A",
	//               "datasource": "Sample Source 1",
	//               "expr": "sample request 1",
	//               "intervalFactor": 0,
	//               "step": 0,
	//               "legendFormat": ""
	//             }
	//           ],
	//           "timeFrom": null,
	//           "timeShift": null,
	//           "tooltip": {
	//             "shared": false,
	//             "value_type": ""
	//           },
	//           "x-axis": true,
	//           "y-axis": true,
	//           "y_formats": null
	//         }
	//       ]
	//     }
	//   ],
	//   "templating": {
	//     "list": null
	//   },
	//   "annotations": {
	//     "list": null
	//   },
	//   "schemiaVersion": 0,
	//   "version": 0,
	//   "links": null,
	//   "time": {
	//     "from": "",
	//     "to": ""
	//   },
	//   "timepicker": {
	//     "refresh_intervals": null,
	//     "time_options": null
	//   }
	// }
}
<commit_msg>Fix unit test.<commit_after>package grafana_test

import (
	"encoding/json"
	"fmt"

	"github.com/grafov/autograf/grafana"
)

func ExampleNewBoard() {
	board := grafana.NewBoard("Sample dashboard title")
	board.ID = 1
	row1 := board.AddRow("Sample row title")
	row1.Add(grafana.NewGraph("Sample graph"))
	graphWithDs := grafana.NewGraph("Sample graph 2")
	target := grafana.Target{
		RefID:      "A",
		Datasource: "Sample Source 1",
		Expr:       "sample request 1"}
	graphWithDs.AddTarget(&target)
	row1.Add(graphWithDs)
	data, _ := json.MarshalIndent(board, "", "  ")
	fmt.Printf("%s", data)
	// Output:
	// {
	//   "id": 1,
	//   "title": "Sample dashboard title",
	//   "originalTitle": "",
	//   "tags": null,
	//   "style": "dark",
	//   "timezone": "browser",
	//   "editable": true,
	//   "hideControls": false,
	//   "sharedCrosshair": false,
	//   "rows": [
	//     {
	//       "title": "Sample row title",
	//       "showTitle": false,
	//       "collapse": false,
	//       "editable": true,
	//       "height": "250px",
	//       "panels": [
	//         {
	//           "id": 1,
	//           "title": "Sample graph",
	//           "span": 12,
	//           "renderer": "flot",
	//           "transparent": false,
	//           "type": "graph",
	//           "error": false,
	//           "isNew": true,
	//           "aliasColors": null,
	//           "bars": false,
	//           "fill": 0,
	//           "grid": {
	//             "leftLogBase": null,
	//             "leftMax": null,
	//             "leftMin": null,
	//             "rightLogBase": null,
	//             "rightMax": null,
	//             "rightMin": null,
	//             "threshold1": null,
	//             "threshold1Color": "",
	//             "threshold2": null,
	//             "threshold2Color": "",
	//             "thresholdLine": false
	//           },
	//           "legend": {
	//             "alignAsTable": false,
	//             "avg": false,
	//             "current": false,
	//             "hideEmpty": false,
	//             "hideZero": false,
	//             "max": false,
	//             "min": false,
	//             "rightSide": false,
	//             "show": false,
	//             "total": false,
	//             "values": false
	//           },
	//           "lines": false,
	//           "linewidth": 0,
	//           "nullPointMode": "connected",
	//           "percentage": false,
	//           "pointradius": 5,
	//           "points": false,
	//           "seriesOverrides": null,
	//           "stack": false,
	//           "steppedLine": false,
	//           "timeFrom": null,
	//           "timeShift": null,
	//           "tooltip": {
	//             "shared": false,
	//             "value_type": ""
	//           },
	//           "x-axis": true,
	//           "y-axis": true,
	//           "y_formats": null
	//         },
	//         {
	//           "id": 2,
	//           "title": "Sample graph 2",
	//           "span": 12,
	//           "renderer": "flot",
	//           "transparent": false,
	//           "type": "graph",
	//           "error": false,
	//           "isNew": true,
	//           "aliasColors": null,
	//           "bars": false,
	//           "fill": 0,
	//           "grid": {
	//             "leftLogBase": null,
	//             "leftMax": null,
	//             "leftMin": null,
	//             "rightLogBase": null,
	//             "rightMax": null,
	//             "rightMin": null,
	//             "threshold1": null,
	//             "threshold1Color": "",
	//             "threshold2": null,
	//             "threshold2Color": "",
	//             "thresholdLine": false
	//           },
	//           "legend": {
	//             "alignAsTable": false,
	//             "avg": false,
	//             "current": false,
	//             "hideEmpty": false,
	//             "hideZero": false,
	//             "max": false,
	//             "min": false,
	//             "rightSide": false,
	//             "show": false,
	//             "total": false,
	//             "values": false
	//           },
	//           "lines": false,
	//           "linewidth": 0,
	//           "nullPointMode": "connected",
	//           "percentage": false,
	//           "pointradius": 5,
	//           "points": false,
	//           "seriesOverrides": null,
	//           "stack": false,
	//           "steppedLine": false,
	//           "targets": [
	//             {
	//               "refId": "A",
	//               "datasource": "Sample Source 1",
	//               "expr": "sample request 1",
	//               "intervalFactor": 0,
	//               "interval": "",
	//               "step": 0,
	//               "legendFormat": ""
	//             }
	//           ],
	//           "timeFrom": null,
	//           "timeShift": null,
	//           "tooltip": {
	//             "shared": false,
	//             "value_type": ""
	//           },
	//           "x-axis": true,
	//           "y-axis": true,
	//           "y_formats": null
	//         }
	//       ]
	//     }
	//   ],
	//   "templating": {
	//     "list": null
	//   },
	//   "annotations": {
	//     "list": null
	//   },
	//   "schemiaVersion": 0,
	//   "version": 0,
	//   "links": null,
	//   "time": {
	//     "from": "",
	//     "to": ""
	//   },
	//   "timepicker": {
	//     "refresh_intervals": null,
	//     "time_options": null
	//   }
	// }
}
<|endoftext|>


<commit_before>package mpb

import (
	"fmt"
	"io"
	"os"
	"sort"
	"sync"
	"time"

	"github.com/vbauerster/mpb/cwriter"
)

type opType uint

const (
	opBarAdd opType = iota
	opBarRemove
)

type SortType uint

const (
	SortNone SortType = iota
	SortTop
	SortBottom
)

// RefreshRate
const rr = 100

// Progress represents the container that renders Progress bars
type Progress struct {
	// WaitGroup for internal rendering sync
	wg *sync.WaitGroup

	out   io.Writer
	width int
	sort  SortType

	op             chan *operation
	rrChangeReqCh  chan time.Duration
	outChangeReqCh chan io.Writer
	countReqCh     chan chan int
	allDone        chan struct{}
}

type operation struct {
	kind   opType
	bar    *Bar
	result chan bool
}

// New returns a new progress bar with defaults
func New() *Progress {
	p := &Progress{
		width:          70,
		op:             make(chan *operation),
		rrChangeReqCh:  make(chan time.Duration),
		outChangeReqCh: make(chan io.Writer),
		countReqCh:     make(chan chan int),
		allDone:        make(chan struct{}),
		wg:             new(sync.WaitGroup),
	}
	go p.server(cwriter.New(os.Stdout), time.NewTicker(rr*time.Millisecond))
	return p
}

// SetWidth sets the width for all underlying bars
func (p *Progress) SetWidth(n int) *Progress {
	if n <= 0 {
		return p
	}
	p.width = n
	return p
}

// SetOut sets underlying writer of progress
// default is os.Stdout
func (p *Progress) SetOut(w io.Writer) *Progress {
	if w == nil {
		return p
	}
	p.outChangeReqCh <- w
	return p
}

// RefreshRate overrides default (30ms) refreshRate value
func (p *Progress) RefreshRate(d time.Duration) *Progress {
	p.rrChangeReqCh <- d
	return p
}

// WithSort sorts the bars, while redering
func (p *Progress) WithSort(sort SortType) *Progress {
	p.sort = sort
	return p
}

// AddBar creates a new progress bar and adds to the container
func (p *Progress) AddBar(total int) *Bar {
	result := make(chan bool)
	bar := newBar(total, p.width, p.wg)
	p.op <- &operation{opBarAdd, bar, result}
	if <-result {
		p.wg.Add(1)
	}
	return bar
}

// RemoveBar removes bar at any time
func (p *Progress) RemoveBar(b *Bar) bool {
	result := make(chan bool)
	p.op <- &operation{opBarRemove, b, result}
	return <-result
}

// BarsCount returns bars count in the container
func (p *Progress) BarsCount() int {
	respCh := make(chan int)
	p.countReqCh <- respCh
	return <-respCh
}

// Stop waits for bars to finish rendering and stops the rendering goroutine
func (p *Progress) Stop() {
	if !p.isAllDone() {
		close(p.allDone)
		p.wg.Wait()
		close(p.op)
	}
}

// server monitors underlying channels and renders any progress bars
func (p *Progress) server(cw *cwriter.Writer, t *time.Ticker) {
	bars := make([]*Bar, 0, 4)
	for {
		select {
		case w := <-p.outChangeReqCh:
			cw.Flush()
			cw = cwriter.New(w)
		case op, ok := <-p.op:
			if !ok {
				t.Stop()
				return
			}
			switch op.kind {
			case opBarAdd:
				bars = append(bars, op.bar)
				op.result <- true
			case opBarRemove:
				var ok bool
				for i, b := range bars {
					if b == op.bar {
						bars = append(bars[:i], bars[i+1:]...)
						ok = true
						b.Stop()
						break
					}
				}
				op.result <- ok
			}
		case respCh := <-p.countReqCh:
			respCh <- len(bars)
		case <-t.C:
			switch p.sort {
			case SortTop:
				sort.Sort(sort.Reverse(SortableBarSlice(bars)))
			case SortBottom:
				sort.Sort(SortableBarSlice(bars))
			}
			for _, b := range bars {
				fmt.Fprintln(cw, b)
			}
			cw.Flush()
			for _, b := range bars {
				go func(b *Bar) {
					b.flushedCh <- struct{}{}
				}(b)
			}
		case d := <-p.rrChangeReqCh:
			t.Stop()
			t = time.NewTicker(d)
		}
	}
}

func (p *Progress) isAllDone() bool {
	select {
	case <-p.allDone:
		return true
	default:
		return false
	}
}
<commit_msg>introduce ErrCallAfterStop error<commit_after>package mpb

import (
	"errors"
	"fmt"
	"io"
	"os"
	"sort"
	"sync"
	"time"

	"github.com/vbauerster/mpb/cwriter"
)

var ErrCallAfterStop = errors.New("method call on stopped Progress instance")

type opType uint

const (
	opBarAdd opType = iota
	opBarRemove
)

type SortType uint

const (
	SortNone SortType = iota
	SortTop
	SortBottom
)

// RefreshRate
const rr = 100

// Progress represents the container that renders Progress bars
type Progress struct {
	// WaitGroup for internal rendering sync
	wg *sync.WaitGroup

	out   io.Writer
	width int
	sort  SortType

	op             chan *operation
	rrChangeReqCh  chan time.Duration
	outChangeReqCh chan io.Writer
	countReqCh     chan chan int
	allDone        chan struct{}
}

type operation struct {
	kind   opType
	bar    *Bar
	result chan bool
}

// New returns a new progress bar with defaults
func New() *Progress {
	p := &Progress{
		width:          70,
		op:             make(chan *operation),
		rrChangeReqCh:  make(chan time.Duration),
		outChangeReqCh: make(chan io.Writer),
		countReqCh:     make(chan chan int),
		allDone:        make(chan struct{}),
		wg:             new(sync.WaitGroup),
	}
	go p.server(cwriter.New(os.Stdout), time.NewTicker(rr*time.Millisecond))
	return p
}

// SetWidth sets the width for all underlying bars
func (p *Progress) SetWidth(n int) *Progress {
	if n <= 0 {
		return p
	}
	p.width = n
	return p
}

// SetOut sets underlying writer of progress. Default is os.Stdout
// pancis, if called on stopped Progress instance, i.e after Stop()
func (p *Progress) SetOut(w io.Writer) *Progress {
	if p.isAllDone() {
		panic(ErrCallAfterStop)
	}
	if w == nil {
		return p
	}
	p.outChangeReqCh <- w
	return p
}

// RefreshRate overrides default (30ms) refreshRate value
// pancis, if called on stopped Progress instance, i.e after Stop()
func (p *Progress) RefreshRate(d time.Duration) *Progress {
	if p.isAllDone() {
		panic(ErrCallAfterStop)
	}
	p.rrChangeReqCh <- d
	return p
}

// WithSort sorts the bars, while redering
func (p *Progress) WithSort(sort SortType) *Progress {
	p.sort = sort
	return p
}

// AddBar creates a new progress bar and adds to the container
// pancis, if called on stopped Progress instance, i.e after Stop()
func (p *Progress) AddBar(total int) *Bar {
	if p.isAllDone() {
		panic(ErrCallAfterStop)
	}
	result := make(chan bool)
	bar := newBar(total, p.width, p.wg)
	p.op <- &operation{opBarAdd, bar, result}
	if <-result {
		p.wg.Add(1)
	}
	return bar
}

// RemoveBar removes bar at any time
// pancis, if called on stopped Progress instance, i.e after Stop()
func (p *Progress) RemoveBar(b *Bar) bool {
	if p.isAllDone() {
		panic(ErrCallAfterStop)
	}
	result := make(chan bool)
	p.op <- &operation{opBarRemove, b, result}
	return <-result
}

// BarsCount returns bars count in the container
// pancis, if called on stopped Progress instance, i.e after Stop()
func (p *Progress) BarsCount() int {
	if p.isAllDone() {
		panic(ErrCallAfterStop)
	}
	respCh := make(chan int)
	p.countReqCh <- respCh
	return <-respCh
}

// Stop waits for bars to finish rendering and stops the rendering goroutine
func (p *Progress) Stop() {
	if !p.isAllDone() {
		close(p.allDone)
		p.wg.Wait()
		close(p.op)
	}
}

// server monitors underlying channels and renders any progress bars
func (p *Progress) server(cw *cwriter.Writer, t *time.Ticker) {
	bars := make([]*Bar, 0, 4)
	for {
		select {
		case w := <-p.outChangeReqCh:
			cw.Flush()
			cw = cwriter.New(w)
		case op, ok := <-p.op:
			if !ok {
				t.Stop()
				return
			}
			switch op.kind {
			case opBarAdd:
				bars = append(bars, op.bar)
				op.result <- true
			case opBarRemove:
				var ok bool
				for i, b := range bars {
					if b == op.bar {
						bars = append(bars[:i], bars[i+1:]...)
						ok = true
						b.Stop()
						break
					}
				}
				op.result <- ok
			}
		case respCh := <-p.countReqCh:
			respCh <- len(bars)
		case <-t.C:
			switch p.sort {
			case SortTop:
				sort.Sort(sort.Reverse(SortableBarSlice(bars)))
			case SortBottom:
				sort.Sort(SortableBarSlice(bars))
			}
			for _, b := range bars {
				fmt.Fprintln(cw, b)
			}
			cw.Flush()
			for _, b := range bars {
				go func(b *Bar) {
					b.flushedCh <- struct{}{}
				}(b)
			}
		case d := <-p.rrChangeReqCh:
			t.Stop()
			t = time.NewTicker(d)
		}
	}
}

func (p *Progress) isAllDone() bool {
	select {
	case <-p.allDone:
		return true
	default:
		return false
	}
}
<|endoftext|>


<commit_before>// puzzle at http://adventofcode.com/day/10

package main

import (
	"bytes"
	"fmt"
	"strconv"
)

var input = "1113122113"
var timesToRepeat = 40
var maxPrintDigits = 60

func main() {

	s := input

	fmt.Printf("%v:\t%v - %v\n", 0, len(s), s)

	for i := 0; i < timesToRepeat; i++ {
		s = NextOutput(s)

		digitsToShow := s
		elipsis := ""
		if len(s) > maxPrintDigits {
			digitsToShow = digitsToShow[:maxPrintDigits]
			elipsis = "..."
		}

		fmt.Printf("%v:\t%v - %v%v\n", i+1, len(s), digitsToShow, elipsis)
	}
}

func NextOutput(s string) string {
	if len(s) == 0 {
		fmt.Println("Warning: Empty input")
		return ""
	}

	var buffer bytes.Buffer

	consecutiveCount := 0
	previousChar := s[0]
	var ch byte

	for i := range s {
		ch = s[i]

		if ch == previousChar {
			consecutiveCount++
		} else {
			buffer.WriteString(strconv.Itoa(consecutiveCount))
			buffer.WriteByte(previousChar)

			consecutiveCount = 1
		}

		previousChar = ch
	}

	buffer.WriteString(strconv.Itoa(consecutiveCount))
	buffer.WriteByte(ch)

	return buffer.String()
}
<commit_msg>Day 10 part 2<commit_after>// puzzle at http://adventofcode.com/day/10

package main

import (
	"bytes"
	"fmt"
	"strconv"
)

var input = "1113122113"
var timesToRepeat = 50
var maxPrintDigits = 60

func main() {

	s := input

	fmt.Printf("%v:\t%v - %v\n", 0, len(s), s)

	for i := 0; i < timesToRepeat; i++ {
		s = NextOutput(s)

		digitsToShow := s
		elipsis := ""
		if len(s) > maxPrintDigits {
			digitsToShow = digitsToShow[:maxPrintDigits]
			elipsis = "..."
		}

		fmt.Printf("%v:\t%v - %v%v\n", i+1, len(s), digitsToShow, elipsis)
	}
}

func NextOutput(s string) string {
	if len(s) == 0 {
		fmt.Println("Warning: Empty input")
		return ""
	}

	var buffer bytes.Buffer

	consecutiveCount := 0
	previousChar := s[0]
	var ch byte

	for i := range s {
		ch = s[i]

		if ch == previousChar {
			consecutiveCount++
		} else {
			buffer.WriteString(strconv.Itoa(consecutiveCount))
			buffer.WriteByte(previousChar)

			consecutiveCount = 1
		}

		previousChar = ch
	}

	buffer.WriteString(strconv.Itoa(consecutiveCount))
	buffer.WriteByte(ch)

	return buffer.String()
}
<|endoftext|>


<commit_before>package gorange

import (
	"fmt"
	"net/http"
	"time"

	gogetter "gopkg.in/karrick/gogetter.v1"
)

const DefaultQueryTimeout = 3 * time.Second

// Querier interface is minimal library abstraction for submitting a query and receiving a response.
type Querier interface {
	Query(string) ([]string, error)
}

// Configurator provides a way to list the range server addresses, and a way to override defaults
// when creating new http.Client instances.
type Configurator struct {
	Addr2Getter   func(string) gogetter.Getter // Addr2Getter converts a range server address to a Getter, ideally a customized http.Client object with a Timeout set. Leave nil to create default gogetter.Getter with DefaultQueryTimeout.
	RetryCallback func(error) bool             // RetryCallback is predicate function that tests whether query should be retried for a given error. Leave nil to retry all errors.
	RetryCount    int                          // RetryCount is number of query retries to be issued if query returns error. Leave 0 to never retry query errors.
	Servers       []string                     // Servers is slice of range server address strings. Must contain at least one string.
	TTL           time.Duration                // TTL is duration of time to cache query responses. Leave 0 to not cache responses.
}

// NewQuerier returns a new instance that sends queries to one or more range servers. The provided
// Configurator not only provides a way of listing one or more range servers, but also allows
// specification of optional retry-on-failure feature and optional TTL cache that memoizes range
// query responses.
//
//    func main() {
//		servers := []string{"range1.example.com", "range2.example.com", "range3.example.com"}
//
//		config := &gorange.Configurator{
//			RetryCount:    len(servers),
//			Servers:       servers,
//			TTL:           5 * time.Minute,
//		}
//
//		// create a range querier; could list additional servers or include other options as well
//		querier, err := gorange.NewQuerier(config)
//		if err != nil {
//			fmt.Fprintf(os.Stderr, "%s", err)
//			os.Exit(1)
//		}
//    }
func NewQuerier(config *Configurator) (Querier, error) {
	if len(config.Servers) == 0 {
		return nil, fmt.Errorf("cannot create Querier without at least one range server address")
	}

	addr2getter := defaultAddr2Getter
	if config.Addr2Getter != nil {
		addr2getter = config.Addr2Getter
	}

	var hg gogetter.Getter

	if len(config.Servers) == 1 {
		hg = addr2getter(config.Servers[0])
	} else {
		rr := &gogetter.RoundRobin{}
		for _, hostname := range config.Servers {
			rr.Getters = append(rr.Getters, addr2getter(hostname))
		}
		hg = rr
	}

	if config.RetryCount > 0 {
		hg = &gogetter.Retrier{
			Getter:        hg,
			RetryCallback: config.RetryCallback,
			RetryCount:    config.RetryCount,
		}
	}

	q := &Client{hg}

	if config.TTL > time.Duration(0) {
		return NewCachingClient(q, config.TTL)
	}

	return q, nil
}

func defaultAddr2Getter(addr string) gogetter.Getter {
	return &gogetter.Prefixer{
		Prefix: fmt.Sprintf("http://%s/range/list?", addr),
		Getter: &http.Client{
			// WARNING: Not having timeout will cause resource leakage if library connects to buggy range server, or a range server over a poor network connection.
			Timeout: time.Duration(DefaultQueryTimeout),

			// Transport: &http.Transport{
			// 	Dial: (&net.Dialer{
			// 		Timeout:   dialTimeout,
			// 		KeepAlive: keepAliveDuration,
			// 	}).Dial,
			// 	MaxIdleConnsPerHost: int(maxConns),
			// },
		},
	}
}
<commit_msg>uses newer gogetter round robin<commit_after>package gorange

import (
	"fmt"
	"net/http"
	"time"

	// gogetter "gopkg.in/karrick/gogetter.v1"
	"github.com/karrick/gogetter"
)

const DefaultQueryTimeout = 3 * time.Second

// Querier interface is minimal library abstraction for submitting a query and receiving a response.
type Querier interface {
	Query(string) ([]string, error)
}

// Configurator provides a way to list the range server addresses, and a way to override defaults
// when creating new http.Client instances.
type Configurator struct {
	// Addr2Getter converts a range server address to a Getter, ideally a customized http.Client
	// object with a Timeout set. Leave nil to create default gogetter.Getter with
	// DefaultQueryTimeout.
	Addr2Getter func(string) gogetter.Getter

	// RetryCallback is predicate function that tests whether query should be retried for a
	// given error. Leave nil to retry all errors.
	RetryCallback func(error) bool

	// RetryCount is number of query retries to be issued if query returns error. Leave 0 to
	// never retry query errors.
	RetryCount int

	// Servers is slice of range server address strings. Must contain at least one string.
	Servers []string

	// TTL is duration of time to cache query responses. Leave 0 to not cache responses.
	TTL time.Duration
}

// NewQuerier returns a new instance that sends queries to one or more range servers. The provided
// Configurator not only provides a way of listing one or more range servers, but also allows
// specification of optional retry-on-failure feature and optional TTL cache that memoizes range
// query responses.
//
//    func main() {
//		servers := []string{"range1.example.com", "range2.example.com", "range3.example.com"}
//
//		config := &gorange.Configurator{
//			RetryCount:    len(servers),
//			Servers:       servers,
//			TTL:           5 * time.Minute,
//		}
//
//		// create a range querier; could list additional servers or include other options as well
//		querier, err := gorange.NewQuerier(config)
//		if err != nil {
//			fmt.Fprintf(os.Stderr, "%s", err)
//			os.Exit(1)
//		}
//    }
func NewQuerier(config *Configurator) (Querier, error) {
	if len(config.Servers) == 0 {
		return nil, fmt.Errorf("cannot create Querier without at least one range server address")
	}

	addr2getter := defaultAddr2Getter
	if config.Addr2Getter != nil {
		addr2getter = config.Addr2Getter
	}

	var hg gogetter.Getter

	if len(config.Servers) == 1 {
		hg = addr2getter(config.Servers[0])
	} else {
		var hostGetters []gogetter.Getter
		for _, hostname := range config.Servers {
			hostGetters = append(hostGetters, addr2getter(hostname))
		}
		hg = gogetter.NewRoundRobin(hostGetters)
	}

	if config.RetryCount > 0 {
		hg = &gogetter.Retrier{
			Getter:        hg,
			RetryCallback: config.RetryCallback,
			RetryCount:    config.RetryCount,
		}
	}

	q := &Client{hg}

	if config.TTL > 0 {
		return NewCachingClient(q, config.TTL)
	}

	return q, nil
}

func defaultAddr2Getter(addr string) gogetter.Getter {
	return &gogetter.Prefixer{
		Prefix: fmt.Sprintf("http://%s/range/list?", addr),
		Getter: &http.Client{
			// WARNING: Not having timeout will cause resource leakage if library connects to buggy range server, or a range server over a poor network connection.
			Timeout: time.Duration(DefaultQueryTimeout),

			// Transport: &http.Transport{
			// 	Dial: (&net.Dialer{
			// 		Timeout:   dialTimeout,
			// 		KeepAlive: keepAliveDuration,
			// 	}).Dial,
			// 	MaxIdleConnsPerHost: int(maxConns),
			// },
		},
	}
}
<|endoftext|>


<commit_before>/* Copyright 2018 Google Inc.

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

    https://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.
*/

// Binary implements a Certificate Management service client.
package main

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"flag"
	"fmt"
	"reflect"
	"time"

	"github.com/google/gnxi/gnoi/cert"
	"github.com/google/gnxi/utils/entity"
	"github.com/kylelemons/godebug/pretty"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	log "github.com/golang/glog"
)

var (
	certID     = flag.String("cert_id", "", "Certificate Management certificate ID.")
	op         = flag.String("op", "get", "Certificate Management operation, one of: provision, install, rotate, get, revoke, check")
	ca         = flag.String("ca", "", "CA certificate file.")
	key        = flag.String("key", "", "Private key file.")
	targetCN   = flag.String("target_name", "", "Common Name of the target.")
	targetAddr = flag.String("target_addr", "localhost:10161", "The target address in the format of host:port")
	timeOut    = flag.Duration("time_out", 5*time.Second, "Timeout for the operation, 5 seconds by default")

	caEnt  *entity.Entity
	ctx    context.Context
	cancel func()
)

func main() {
	flag.Parse()

	if *ca == "" || *key == "" {
		log.Exit("-ca and -key must be set with file locations")
	}
	if *targetCN == "" {
		log.Exit("Must set a Common Name ID with -targetCN.")
	}

	var err error
	if caEnt, err = entity.FromFile(*ca, *key); err != nil {
		log.Exitf("Failed to load certificate and key from file: %v", err)
	}

	ctx, cancel = context.WithTimeout(context.Background(), *timeOut)
	defer cancel()

	switch *op {
	case "provision":
		certIDCheck()
		provision()
		break
	case "install":
		certIDCheck()
		install()
		break
	case "rotate":
		certIDCheck()
		rotate()
		break
	case "revoke":
		revoke()
		break
	case "check":
		check()
		break
	case "get":
		get()
		break
	default:
		log.Exitf("Unknown operation: %q", *op)
	}
}

func certIDCheck() {
	if *certID == "" {
		log.Exit("Must set a certificate ID with -cert_id.")
	}
}

// gnoiEncrypted creates an encrypted TLS connection to the target.
func gnoiEncrypted(c tls.Certificate) (*grpc.ClientConn, *cert.Client) {
	opts := []grpc.DialOption{grpc.WithTransportCredentials(credentials.NewTLS(
		&tls.Config{
			InsecureSkipVerify: true,
			Certificates:       []tls.Certificate{c},
			RootCAs:            nil,
		}))}

	conn, err := grpc.Dial(*targetAddr, opts...)
	if err != nil {
		log.Exitf("Failed dial to %q: %v", *targetAddr, err)
	}

	client := cert.NewClient(conn)
	return conn, client
}

// gnoiAuthenticated creates an authenticated TLS connection to the target.
func gnoiAuthenticated(targetName string) (*grpc.ClientConn, *cert.Client) {
	clientEnt, err := entity.CreateSigned("client", nil, caEnt)
	if err != nil {
		log.Exitf("Failed to create a signed entity: %v", err)
	}
	caPool := x509.NewCertPool()
	caPool.AddCert(caEnt.Certificate.Leaf)

	opts := []grpc.DialOption{grpc.WithTransportCredentials(credentials.NewTLS(
		&tls.Config{
			ServerName:   targetName,
			Certificates: []tls.Certificate{*clientEnt.Certificate},
			RootCAs:      caPool,
		}))}

	conn, err := grpc.Dial(*targetAddr, opts...)
	if err != nil {
		log.Exitf("Failed dial to %q: %v", *targetAddr, err)
	}

	client := cert.NewClient(conn)
	return conn, client
}

// signer is called to create a Certificate from a CSR.
func signer(csr *x509.CertificateRequest) (*x509.Certificate, error) {
	e, err := entity.FromSigningRequest(csr)
	if err != nil {
		return nil, fmt.Errorf("failed generating a cert from a CSR: %v", err)
	}
	if err := e.SignWith(caEnt); err != nil {
		return nil, fmt.Errorf("failed to sign the certificate: %v", err)
	}
	return e.Certificate.Leaf, nil
}

// provision provisions a target in bootstrapping mode.
func provision() {
	// Using the CA x509 cert as default Certificate, but can be any.
	conn, client := gnoiEncrypted(*caEnt.Certificate)
	defer conn.Close()

	if err := client.Install(ctx, *certID, pkix.Name{CommonName: *targetCN}, signer, []*x509.Certificate{caEnt.Certificate.Leaf}); err != nil {
		log.Exit("Failed Install:", err)
	}
	log.Info("Install success")
}

// install installs a certificate in authenticated mode.
func install() {
	conn, client := gnoiAuthenticated(*targetCN)
	defer conn.Close()

	if err := client.Install(ctx, *certID, pkix.Name{CommonName: *targetCN}, signer, []*x509.Certificate{caEnt.Certificate.Leaf}); err != nil {
		log.Exit("Failed Install:", err)
	}
	log.Info("Install success")
}

// rotate rotates a certificate in authenticated mode.
func rotate() {
	conn, client := gnoiAuthenticated(*targetCN)
	defer conn.Close()

	if err := client.Rotate(ctx, *certID, pkix.Name{CommonName: *targetCN}, signer, []*x509.Certificate{caEnt.Certificate.Leaf}, func() error { return nil }); err != nil {
		log.Exit("Failed Rotate:", err)
	}
	log.Info("Rotate success")
}

// revoke revokes a certificate in authenticated mode.
func revoke() {
	if *certID == "" {
		log.Exit("Must set a certificate ID with -cert_id.")
	}
	conn, client := gnoiAuthenticated(*targetCN)
	defer conn.Close()

	revoked, err := client.RevokeCertificates(ctx, []string{*certID})
	if err != nil {
		log.Exit("Failed RevokeCertificates:", err)
	}
	log.Info("RevokeCertificates:\n", pretty.Sprint(revoked))
}

// revoke checks if a target can generate certificates - authenticated mode.
func check() {
	conn, client := gnoiAuthenticated(*targetCN)
	defer conn.Close()

	resp, err := client.CanGenerateCSR(ctx)
	if err != nil {
		log.Exit("Failed CanGenerateCSR:", err)
	}
	log.Info("CanGenerateCSR:\n", pretty.Sprint(resp))
}

// get fetches the installed certificates on a target - authenticated mode.
func get() {
	conn, client := gnoiAuthenticated(*targetCN)
	defer conn.Close()

	resp, err := client.GetCertificates(ctx)
	if err != nil {
		log.Exit("Failed GetCertificates:", err)
	}

	pretty.DefaultFormatter[reflect.TypeOf(&x509.Certificate{})] = func(c *x509.Certificate) string {
		return pretty.Sprint(c.Subject.CommonName)
	}
	log.Info("GetCertificates:\n", pretty.Sprint(resp))
}
<commit_msg>Remove redundant breaks.<commit_after>/* Copyright 2018 Google Inc.

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

    https://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.
*/

// Binary implements a Certificate Management service client.
package main

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"flag"
	"fmt"
	"reflect"
	"time"

	"github.com/google/gnxi/gnoi/cert"
	"github.com/google/gnxi/utils/entity"
	"github.com/kylelemons/godebug/pretty"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	log "github.com/golang/glog"
)

var (
	certID     = flag.String("cert_id", "", "Certificate Management certificate ID.")
	op         = flag.String("op", "get", "Certificate Management operation, one of: provision, install, rotate, get, revoke, check")
	ca         = flag.String("ca", "", "CA certificate file.")
	key        = flag.String("key", "", "Private key file.")
	targetCN   = flag.String("target_name", "", "Common Name of the target.")
	targetAddr = flag.String("target_addr", "localhost:10161", "The target address in the format of host:port")
	timeOut    = flag.Duration("time_out", 5*time.Second, "Timeout for the operation, 5 seconds by default")

	caEnt  *entity.Entity
	ctx    context.Context
	cancel func()
)

func main() {
	flag.Parse()

	if *ca == "" || *key == "" {
		log.Exit("-ca and -key must be set with file locations")
	}
	if *targetCN == "" {
		log.Exit("Must set a Common Name ID with -targetCN.")
	}

	var err error
	if caEnt, err = entity.FromFile(*ca, *key); err != nil {
		log.Exitf("Failed to load certificate and key from file: %v", err)
	}

	ctx, cancel = context.WithTimeout(context.Background(), *timeOut)
	defer cancel()

	switch *op {
	case "provision":
		certIDCheck()
		provision()
	case "install":
		certIDCheck()
		install()
	case "rotate":
		certIDCheck()
		rotate()
	case "revoke":
		revoke()
	case "check":
		check()
	case "get":
		get()
	default:
		log.Exitf("Unknown operation: %q", *op)
	}
}

func certIDCheck() {
	if *certID == "" {
		log.Exit("Must set a certificate ID with -cert_id.")
	}
}

// gnoiEncrypted creates an encrypted TLS connection to the target.
func gnoiEncrypted(c tls.Certificate) (*grpc.ClientConn, *cert.Client) {
	opts := []grpc.DialOption{grpc.WithTransportCredentials(credentials.NewTLS(
		&tls.Config{
			InsecureSkipVerify: true,
			Certificates:       []tls.Certificate{c},
			RootCAs:            nil,
		}))}

	conn, err := grpc.Dial(*targetAddr, opts...)
	if err != nil {
		log.Exitf("Failed dial to %q: %v", *targetAddr, err)
	}

	client := cert.NewClient(conn)
	return conn, client
}

// gnoiAuthenticated creates an authenticated TLS connection to the target.
func gnoiAuthenticated(targetName string) (*grpc.ClientConn, *cert.Client) {
	clientEnt, err := entity.CreateSigned("client", nil, caEnt)
	if err != nil {
		log.Exitf("Failed to create a signed entity: %v", err)
	}
	caPool := x509.NewCertPool()
	caPool.AddCert(caEnt.Certificate.Leaf)

	opts := []grpc.DialOption{grpc.WithTransportCredentials(credentials.NewTLS(
		&tls.Config{
			ServerName:   targetName,
			Certificates: []tls.Certificate{*clientEnt.Certificate},
			RootCAs:      caPool,
		}))}

	conn, err := grpc.Dial(*targetAddr, opts...)
	if err != nil {
		log.Exitf("Failed dial to %q: %v", *targetAddr, err)
	}

	client := cert.NewClient(conn)
	return conn, client
}

// signer is called to create a Certificate from a CSR.
func signer(csr *x509.CertificateRequest) (*x509.Certificate, error) {
	e, err := entity.FromSigningRequest(csr)
	if err != nil {
		return nil, fmt.Errorf("failed generating a cert from a CSR: %v", err)
	}
	if err := e.SignWith(caEnt); err != nil {
		return nil, fmt.Errorf("failed to sign the certificate: %v", err)
	}
	return e.Certificate.Leaf, nil
}

// provision provisions a target in bootstrapping mode.
func provision() {
	// Using the CA x509 cert as default Certificate, but can be any.
	conn, client := gnoiEncrypted(*caEnt.Certificate)
	defer conn.Close()

	if err := client.Install(ctx, *certID, pkix.Name{CommonName: *targetCN}, signer, []*x509.Certificate{caEnt.Certificate.Leaf}); err != nil {
		log.Exit("Failed Install:", err)
	}
	log.Info("Install success")
}

// install installs a certificate in authenticated mode.
func install() {
	conn, client := gnoiAuthenticated(*targetCN)
	defer conn.Close()

	if err := client.Install(ctx, *certID, pkix.Name{CommonName: *targetCN}, signer, []*x509.Certificate{caEnt.Certificate.Leaf}); err != nil {
		log.Exit("Failed Install:", err)
	}
	log.Info("Install success")
}

// rotate rotates a certificate in authenticated mode.
func rotate() {
	conn, client := gnoiAuthenticated(*targetCN)
	defer conn.Close()

	if err := client.Rotate(ctx, *certID, pkix.Name{CommonName: *targetCN}, signer, []*x509.Certificate{caEnt.Certificate.Leaf}, func() error { return nil }); err != nil {
		log.Exit("Failed Rotate:", err)
	}
	log.Info("Rotate success")
}

// revoke revokes a certificate in authenticated mode.
func revoke() {
	if *certID == "" {
		log.Exit("Must set a certificate ID with -cert_id.")
	}
	conn, client := gnoiAuthenticated(*targetCN)
	defer conn.Close()

	revoked, err := client.RevokeCertificates(ctx, []string{*certID})
	if err != nil {
		log.Exit("Failed RevokeCertificates:", err)
	}
	log.Info("RevokeCertificates:\n", pretty.Sprint(revoked))
}

// revoke checks if a target can generate certificates - authenticated mode.
func check() {
	conn, client := gnoiAuthenticated(*targetCN)
	defer conn.Close()

	resp, err := client.CanGenerateCSR(ctx)
	if err != nil {
		log.Exit("Failed CanGenerateCSR:", err)
	}
	log.Info("CanGenerateCSR:\n", pretty.Sprint(resp))
}

// get fetches the installed certificates on a target - authenticated mode.
func get() {
	conn, client := gnoiAuthenticated(*targetCN)
	defer conn.Close()

	resp, err := client.GetCertificates(ctx)
	if err != nil {
		log.Exit("Failed GetCertificates:", err)
	}

	pretty.DefaultFormatter[reflect.TypeOf(&x509.Certificate{})] = func(c *x509.Certificate) string {
		return pretty.Sprint(c.Subject.CommonName)
	}
	log.Info("GetCertificates:\n", pretty.Sprint(resp))
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"math"
	"os"
	"strconv"
)

const perPage int = 30

func calcRealPage(cutRange, page int) int {
	maxPos := cutRange + page + 2
	return int(math.Ceil(float64(maxPos) / float64(perPage)))
}

func calcRange(cutRange, page int) (int, int) {
	start := 0
	end := (cutRange+page+1)%perPage + 1
	if page > 1 {
		start = end - 1
	}
	return start, end
}

func createPosts(page int) []string {
	posts := make([]string, 0)
	for i := 0; i < perPage; i++ {
		n := (page-1)*perPage + i + 1
		if n > 100 {
			break
		}
		fuck := fmt.Sprintf("Posts%03d", n)
		posts = append(posts, fuck)
	}
	return posts
}

func genPages(cutRange, pageStart, pageEnd int) {
	for page := pageStart; page <= pageEnd; page++ {
		realPage := calcRealPage(cutRange, page)
		posts := createPosts(realPage)
		start, end := calcRange(cutRange, page)
		if end > len(posts) {
			fmt.Printf("Page %02d: []\n", page)
		} else {
			fmt.Printf("Page %02d: %s\n", page, posts[start:end])
		}
	}
}

func main() {
	args := os.Args[1:]
	if len(args) == 0 {
		log.Fatal("Bad arguments!")
	}
	cutRange, _ := strconv.Atoi(args[0])
	pageStart, _ := strconv.Atoi(args[1])
	pageEnd, _ := strconv.Atoi(args[2])
	genPages(cutRange, pageStart, pageEnd)
}
<commit_msg>placeholder: a memoize function in Go<commit_after>package main

import (
	"fmt"
	"log"
	"math"
	"os"
	"strconv"
)

const perPage int = 30

type createPostFunc func(int) []string

func memoize(f createPostFunc) createPostFunc {
	cache := make(map[int][]string)
	return func(page int) []string {
		if posts, ok := cache[page]; ok {
			return posts
		}
		cache[page] = f(page)
		return cache[page]
	}
}

func calcRealPage(cutRange, page int) int {
	maxPos := cutRange + page + 2
	return int(math.Ceil(float64(maxPos) / float64(perPage)))
}

func calcRange(cutRange, page int) (int, int) {
	start := 0
	end := (cutRange+page+1)%perPage + 1
	if page > 1 {
		start = end - 1
	}
	return start, end
}

func createPosts(page int) []string {
	posts := make([]string, 0)
	for i := 0; i < perPage; i++ {
		n := (page-1)*perPage + i + 1
		if n > 100 {
			break
		}
		fuck := fmt.Sprintf("Posts%03d", n)
		posts = append(posts, fuck)
	}
	return posts
}

func genPages(cutRange, pageStart, pageEnd int) {
	createPostsCache := memoize(createPosts)
	for page := pageStart; page <= pageEnd; page++ {
		realPage := calcRealPage(cutRange, page)
		posts := createPostsCache(realPage)
		start, end := calcRange(cutRange, page)
		if end > len(posts) {
			fmt.Printf("Page %02d: []\n", page)
		} else {
			fmt.Printf("Page %02d: %s\n", page, posts[start:end])
		}
	}
}

func main() {
	args := os.Args[1:]
	if len(args) == 0 {
		log.Fatal("Bad arguments!")
	}
	cutRange, _ := strconv.Atoi(args[0])
	pageStart, _ := strconv.Atoi(args[1])
	pageEnd, _ := strconv.Atoi(args[2])
	genPages(cutRange, pageStart, pageEnd)
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC
//
// 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 keys

import (
	"errors"
	"reflect"
	"testing"

	"github.com/google/chrome-ssh-agent/go/chrome/fakes"
	"github.com/gopherjs/gopherjs/js"
	"github.com/kr/pretty"
)

type dummyManager struct {
	ID             ID
	Name           string
	PEMPrivateKey  string
	Passphrase     string
	ConfiguredKeys []*ConfiguredKey
	LoadedKeys     []*LoadedKey
	Key            *LoadedKey
	Err            error
}

func (m *dummyManager) Configured(callback func(keys []*ConfiguredKey, err error)) {
	callback(m.ConfiguredKeys, m.Err)
}

func (m *dummyManager) Add(name string, pemPrivateKey string, callback func(err error)) {
	m.Name = name
	m.PEMPrivateKey = pemPrivateKey
	callback(m.Err)
}

func (m *dummyManager) Remove(id ID, callback func(err error)) {
	m.ID = id
	callback(m.Err)
}

func (m *dummyManager) Loaded(callback func(keys []*LoadedKey, err error)) {
	callback(m.LoadedKeys, m.Err)
}

func (m *dummyManager) Load(id ID, passphrase string, callback func(err error)) {
	m.ID = id
	m.Passphrase = passphrase
	callback(m.Err)
}

func (m *dummyManager) Unload(key *LoadedKey, callback func(err error)) {
	m.Key = key
	callback(m.Err)
}

func TestClientServerConfigured(t *testing.T) {
	hub := fakes.NewMessageHub()
	mgr := &dummyManager{}
	cli := NewClient(hub)
	NewServer(mgr, hub)

	k0 := &ConfiguredKey{Object: js.Global.Get("Object").New()}
	k0.ID = ID("id-0")
	k0.Name = "key-0"
	k1 := &ConfiguredKey{Object: js.Global.Get("Object").New()}
	k1.ID = ID("id-1")
	k1.Name = "key-1"

	wantConfiguredKeys := []*ConfiguredKey{k0, k1}
	wantErr := errors.New("failed")

	mgr.ConfiguredKeys = append(mgr.ConfiguredKeys, wantConfiguredKeys...)
	mgr.Err = wantErr

	configured, err := syncConfigured(cli)
	// Compare using reflect.DeepEqual since pretty.Diff fails to
	// terminate on this input.
	if !reflect.DeepEqual(configured, wantConfiguredKeys) {
		t.Errorf("incorrect configured keys; got %s, want %s", configured, wantConfiguredKeys)
	}
	if diff := pretty.Diff(err, wantErr); diff != nil {
		t.Errorf("incorrect error; -got +want: %s", diff)
	}
}

func TestClientServerAdd(t *testing.T) {
	hub := fakes.NewMessageHub()
	mgr := &dummyManager{}
	cli := NewClient(hub)
	NewServer(mgr, hub)

	wantName := "some-name"
	wantPrivateKey := "private-key"
	wantErr := errors.New("failed")

	mgr.Err = wantErr

	err := syncAdd(cli, wantName, wantPrivateKey)
	if diff := pretty.Diff(mgr.Name, wantName); diff != nil {
		t.Errorf("incorrect name; -got +want: %s", diff)
	}
	if diff := pretty.Diff(mgr.PEMPrivateKey, wantPrivateKey); diff != nil {
		t.Errorf("incorrect private key; -got +want: %s", diff)
	}
	if diff := pretty.Diff(err, wantErr); diff != nil {
		t.Errorf("incorrect error; -got +want: %s", diff)
	}
}

func TestClientServerRemove(t *testing.T) {
	hub := fakes.NewMessageHub()
	mgr := &dummyManager{}
	cli := NewClient(hub)
	NewServer(mgr, hub)

	wantID := ID("id-0")
	wantErr := errors.New("failed")

	mgr.Err = wantErr

	err := syncRemove(cli, wantID)
	if diff := pretty.Diff(mgr.ID, wantID); diff != nil {
		t.Errorf("incorrect ID; -got +want: %s", diff)
	}
	if diff := pretty.Diff(err, wantErr); diff != nil {
		t.Errorf("incorrect error; -got +want: %s", diff)
	}
}

func TestClientServerLoaded(t *testing.T) {
	hub := fakes.NewMessageHub()
	mgr := &dummyManager{}
	cli := NewClient(hub)
	NewServer(mgr, hub)

	k0 := &LoadedKey{Object: js.Global.Get("Object").New()}
	k0.Type = "type-0"
	k0.SetBlob([]byte("blob-0"))
	k0.Comment = "comment-0"
	k1 := &LoadedKey{Object: js.Global.Get("Object").New()}
	k1.Type = "type-1"
	k1.SetBlob([]byte("blob-1"))
	k1.Comment = "comment-1"

	wantLoadedKeys := []*LoadedKey{k0, k1}
	wantErr := errors.New("failed")

	mgr.LoadedKeys = append(mgr.LoadedKeys, wantLoadedKeys...)
	mgr.Err = wantErr

	loaded, err := syncLoaded(cli)
	// Compare using reflect.DeepEqual since pretty.Diff fails to
	// terminate on this input.
	if !reflect.DeepEqual(loaded, wantLoadedKeys) {
		t.Errorf("incorrect loaded keys; got %s, want %s", loaded, wantLoadedKeys)
	}
	if diff := pretty.Diff(err, wantErr); diff != nil {
		t.Errorf("incorrect error; -got +want: %s", diff)
	}
}

func TestClientServerLoad(t *testing.T) {
	hub := fakes.NewMessageHub()
	mgr := &dummyManager{}
	cli := NewClient(hub)
	NewServer(mgr, hub)

	wantID := ID("id-0")
	wantPassphrase := "secret"
	wantErr := errors.New("failed")

	mgr.Err = wantErr

	err := syncLoad(cli, wantID, wantPassphrase)
	if diff := pretty.Diff(mgr.ID, wantID); diff != nil {
		t.Errorf("incorrect ID; -got +want: %s", diff)
	}
	if diff := pretty.Diff(mgr.Passphrase, wantPassphrase); diff != nil {
		t.Errorf("incorrect passphrase; -got +want: %s", diff)
	}
	if diff := pretty.Diff(err, wantErr); diff != nil {
		t.Errorf("incorrect error; -got +want: %s", diff)
	}
}

func TestClientServerUnload(t *testing.T) {
	hub := fakes.NewMessageHub()
	mgr := &dummyManager{}
	cli := NewClient(hub)
	NewServer(mgr, hub)

	wantKey := &LoadedKey{Object: js.Global.Get("Object").New()}
	wantKey.Type = "type-0"
	wantKey.SetBlob([]byte("blob-0"))
	wantKey.Comment = "comment1"
	wantErr := errors.New("failed")

	mgr.Err = wantErr

	err := syncUnload(cli, wantKey)
	// Compare using reflect.DeepEqual since pretty.Diff causes test to fail
	// without any output.
	if !reflect.DeepEqual(mgr.Key, wantKey) {
		t.Errorf("incorrect key; got %s, want %s", mgr.Key, wantKey)
	}
	if diff := pretty.Diff(err, wantErr); diff != nil {
		t.Errorf("incorrect error; -got +want: %s", diff)
	}
}
<commit_msg>Fix printf args to satisfy go's vetting<commit_after>// Copyright 2018 Google LLC
//
// 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 keys

import (
	"errors"
	"reflect"
	"testing"

	"github.com/google/chrome-ssh-agent/go/chrome/fakes"
	"github.com/gopherjs/gopherjs/js"
	"github.com/kr/pretty"
)

type dummyManager struct {
	ID             ID
	Name           string
	PEMPrivateKey  string
	Passphrase     string
	ConfiguredKeys []*ConfiguredKey
	LoadedKeys     []*LoadedKey
	Key            *LoadedKey
	Err            error
}

func (m *dummyManager) Configured(callback func(keys []*ConfiguredKey, err error)) {
	callback(m.ConfiguredKeys, m.Err)
}

func (m *dummyManager) Add(name string, pemPrivateKey string, callback func(err error)) {
	m.Name = name
	m.PEMPrivateKey = pemPrivateKey
	callback(m.Err)
}

func (m *dummyManager) Remove(id ID, callback func(err error)) {
	m.ID = id
	callback(m.Err)
}

func (m *dummyManager) Loaded(callback func(keys []*LoadedKey, err error)) {
	callback(m.LoadedKeys, m.Err)
}

func (m *dummyManager) Load(id ID, passphrase string, callback func(err error)) {
	m.ID = id
	m.Passphrase = passphrase
	callback(m.Err)
}

func (m *dummyManager) Unload(key *LoadedKey, callback func(err error)) {
	m.Key = key
	callback(m.Err)
}

func TestClientServerConfigured(t *testing.T) {
	hub := fakes.NewMessageHub()
	mgr := &dummyManager{}
	cli := NewClient(hub)
	NewServer(mgr, hub)

	k0 := &ConfiguredKey{Object: js.Global.Get("Object").New()}
	k0.ID = ID("id-0")
	k0.Name = "key-0"
	k1 := &ConfiguredKey{Object: js.Global.Get("Object").New()}
	k1.ID = ID("id-1")
	k1.Name = "key-1"

	wantConfiguredKeys := []*ConfiguredKey{k0, k1}
	wantErr := errors.New("failed")

	mgr.ConfiguredKeys = append(mgr.ConfiguredKeys, wantConfiguredKeys...)
	mgr.Err = wantErr

	configured, err := syncConfigured(cli)
	// Compare using reflect.DeepEqual since pretty.Diff fails to
	// terminate on this input.
	if !reflect.DeepEqual(configured, wantConfiguredKeys) {
		t.Errorf("incorrect configured keys; got %v, want %v", configured, wantConfiguredKeys)
	}
	if diff := pretty.Diff(err, wantErr); diff != nil {
		t.Errorf("incorrect error; -got +want: %s", diff)
	}
}

func TestClientServerAdd(t *testing.T) {
	hub := fakes.NewMessageHub()
	mgr := &dummyManager{}
	cli := NewClient(hub)
	NewServer(mgr, hub)

	wantName := "some-name"
	wantPrivateKey := "private-key"
	wantErr := errors.New("failed")

	mgr.Err = wantErr

	err := syncAdd(cli, wantName, wantPrivateKey)
	if diff := pretty.Diff(mgr.Name, wantName); diff != nil {
		t.Errorf("incorrect name; -got +want: %s", diff)
	}
	if diff := pretty.Diff(mgr.PEMPrivateKey, wantPrivateKey); diff != nil {
		t.Errorf("incorrect private key; -got +want: %s", diff)
	}
	if diff := pretty.Diff(err, wantErr); diff != nil {
		t.Errorf("incorrect error; -got +want: %s", diff)
	}
}

func TestClientServerRemove(t *testing.T) {
	hub := fakes.NewMessageHub()
	mgr := &dummyManager{}
	cli := NewClient(hub)
	NewServer(mgr, hub)

	wantID := ID("id-0")
	wantErr := errors.New("failed")

	mgr.Err = wantErr

	err := syncRemove(cli, wantID)
	if diff := pretty.Diff(mgr.ID, wantID); diff != nil {
		t.Errorf("incorrect ID; -got +want: %s", diff)
	}
	if diff := pretty.Diff(err, wantErr); diff != nil {
		t.Errorf("incorrect error; -got +want: %s", diff)
	}
}

func TestClientServerLoaded(t *testing.T) {
	hub := fakes.NewMessageHub()
	mgr := &dummyManager{}
	cli := NewClient(hub)
	NewServer(mgr, hub)

	k0 := &LoadedKey{Object: js.Global.Get("Object").New()}
	k0.Type = "type-0"
	k0.SetBlob([]byte("blob-0"))
	k0.Comment = "comment-0"
	k1 := &LoadedKey{Object: js.Global.Get("Object").New()}
	k1.Type = "type-1"
	k1.SetBlob([]byte("blob-1"))
	k1.Comment = "comment-1"

	wantLoadedKeys := []*LoadedKey{k0, k1}
	wantErr := errors.New("failed")

	mgr.LoadedKeys = append(mgr.LoadedKeys, wantLoadedKeys...)
	mgr.Err = wantErr

	loaded, err := syncLoaded(cli)
	// Compare using reflect.DeepEqual since pretty.Diff fails to
	// terminate on this input.
	if !reflect.DeepEqual(loaded, wantLoadedKeys) {
		t.Errorf("incorrect loaded keys; got %s, want %s", loaded, wantLoadedKeys)
	}
	if diff := pretty.Diff(err, wantErr); diff != nil {
		t.Errorf("incorrect error; -got +want: %s", diff)
	}
}

func TestClientServerLoad(t *testing.T) {
	hub := fakes.NewMessageHub()
	mgr := &dummyManager{}
	cli := NewClient(hub)
	NewServer(mgr, hub)

	wantID := ID("id-0")
	wantPassphrase := "secret"
	wantErr := errors.New("failed")

	mgr.Err = wantErr

	err := syncLoad(cli, wantID, wantPassphrase)
	if diff := pretty.Diff(mgr.ID, wantID); diff != nil {
		t.Errorf("incorrect ID; -got +want: %s", diff)
	}
	if diff := pretty.Diff(mgr.Passphrase, wantPassphrase); diff != nil {
		t.Errorf("incorrect passphrase; -got +want: %s", diff)
	}
	if diff := pretty.Diff(err, wantErr); diff != nil {
		t.Errorf("incorrect error; -got +want: %s", diff)
	}
}

func TestClientServerUnload(t *testing.T) {
	hub := fakes.NewMessageHub()
	mgr := &dummyManager{}
	cli := NewClient(hub)
	NewServer(mgr, hub)

	wantKey := &LoadedKey{Object: js.Global.Get("Object").New()}
	wantKey.Type = "type-0"
	wantKey.SetBlob([]byte("blob-0"))
	wantKey.Comment = "comment1"
	wantErr := errors.New("failed")

	mgr.Err = wantErr

	err := syncUnload(cli, wantKey)
	// Compare using reflect.DeepEqual since pretty.Diff causes test to fail
	// without any output.
	if !reflect.DeepEqual(mgr.Key, wantKey) {
		t.Errorf("incorrect key; got %s, want %s", mgr.Key, wantKey)
	}
	if diff := pretty.Diff(err, wantErr); diff != nil {
		t.Errorf("incorrect error; -got +want: %s", diff)
	}
}
<|endoftext|>


<commit_before>package models

import (
	"bytes"
	"compress/gzip"
	"encoding/binary"
	"github.com/lunixbochs/struc"
)

// savestate format:
// https://github.com/lunixbochs/usercorn/issues/176

// file header
// uint32(savestate format version)
// -- unicorn header --
// uint32(unicorn major version)
// uint32(unicorn minor version)
// uint32(unicorn arch enum)
// uint32(unicorn mode enum)
//
// -- compressed data header --
// uint64(length of compressed data)
// remainder is gzip-compressed
//
// -- uncompressed data start --
// registers
// uint32(number of registers)
// 1..num: uint32(register enum), uint64(register value)
//
// memory
// uint64(number of mapped sections)
// 1..num: uint64(addr), uint64(len), uint32(prot), <raw memory bytes of len>

type SaveHeader struct {
	Version          uint32
	UcMajor, UcMinor uint32
	UcArch, UcMode   uint32

	BodySize   uint64 `struc:"sizeof=Compressed"`
	Compressed []byte
}

func (s *SaveHeader) PackBody(b *SaveBody) error {
	var tmp bytes.Buffer
	gz := gzip.NewWriter(&tmp)
	err := struc.PackWithOptions(gz, b, &struc.Options{Order: binary.BigEndian})
	if err != nil {
		return err
	}
	s.Compressed = tmp.Bytes()
	return nil
}

func (s *SaveHeader) UnpackBody() (*SaveBody, error) {
	gz, err := gzip.NewReader(bytes.NewReader(s.Compressed))
	if err != nil {
		return nil, err
	}
	body := &SaveBody{}
	err = struc.UnpackWithOptions(gz, body, &struc.Options{Order: binary.BigEndian})
	if err != nil {
		return nil, err
	}
	return body, nil
}

type SaveReg struct {
	Enum, Val uint64
}

type SaveMem struct {
	Addr, Size uint64
	Prot       uint32

	Len  uint64 `struc:"sizeof=Data"`
	Data []byte
}

type SaveBody struct {
	RegCount uint64 `struc:"sizeof=Regs"`
	Regs     []SaveReg
	MemCount uint64 `struc:"sizeof=Mem"`
	Mem      []SaveMem
}

// TODO: pack using all structs above instead of just header
func Save(u Usercorn) ([]byte, error) {
	var buf bytes.Buffer
	arch := u.Arch()
	options := &struc.Options{Order: binary.BigEndian}
	// build compressed body
	s := StrucStream{&buf, options}

	// register list
	s.Pack(uint64(len(arch.Regs)))
	for _, enum := range arch.Regs {
		val, _ := u.RegRead(enum)
		s.Pack(uint64(enum), uint64(val))
	}

	// memory mappings
	mappings := u.Mappings()
	s.Pack(uint64(len(mappings)))
	for _, m := range mappings {
		s.Pack(uint64(m.Addr), uint64(m.Size), uint32(m.Prot))
		mem, err := u.MemRead(m.Addr, m.Size)
		if err != nil {
			return nil, err
		}
		buf.Write(mem)
	}

	// compress body
	var tmp bytes.Buffer
	gz := gzip.NewWriter(&tmp)
	buf.WriteTo(gz)
	buf.Reset()

	// write header / combine everything
	header := &SaveHeader{
		Version: 1,
		// unicorn version isn't exposed by Go bindings yet (Unicorn PR #483)
		UcMajor: 0, UcMinor: 0,
		UcArch: uint32(arch.UC_ARCH), UcMode: uint32(arch.UC_MODE),
		Compressed: tmp.Bytes(),
	}
	var final bytes.Buffer
	struc.PackWithOptions(&final, header, options)
	return final.Bytes(), nil
}
<commit_msg>skip memory errors during state save<commit_after>package models

import (
	"bytes"
	"compress/gzip"
	"encoding/binary"
	"fmt"
	"github.com/lunixbochs/struc"
	"os"
)

// savestate format:
// https://github.com/lunixbochs/usercorn/issues/176

// file header
// uint32(savestate format version)
// -- unicorn header --
// uint32(unicorn major version)
// uint32(unicorn minor version)
// uint32(unicorn arch enum)
// uint32(unicorn mode enum)
//
// -- compressed data header --
// uint64(length of compressed data)
// remainder is gzip-compressed
//
// -- uncompressed data start --
// registers
// uint32(number of registers)
// 1..num: uint32(register enum), uint64(register value)
//
// memory
// uint64(number of mapped sections)
// 1..num: uint64(addr), uint64(len), uint32(prot), <raw memory bytes of len>

type SaveHeader struct {
	Version          uint32
	UcMajor, UcMinor uint32
	UcArch, UcMode   uint32

	BodySize   uint64 `struc:"sizeof=Compressed"`
	Compressed []byte
}

func (s *SaveHeader) PackBody(b *SaveBody) error {
	var tmp bytes.Buffer
	gz := gzip.NewWriter(&tmp)
	err := struc.PackWithOptions(gz, b, &struc.Options{Order: binary.BigEndian})
	if err != nil {
		return err
	}
	s.Compressed = tmp.Bytes()
	return nil
}

func (s *SaveHeader) UnpackBody() (*SaveBody, error) {
	gz, err := gzip.NewReader(bytes.NewReader(s.Compressed))
	if err != nil {
		return nil, err
	}
	body := &SaveBody{}
	err = struc.UnpackWithOptions(gz, body, &struc.Options{Order: binary.BigEndian})
	if err != nil {
		return nil, err
	}
	return body, nil
}

type SaveReg struct {
	Enum, Val uint64
}

type SaveMem struct {
	Addr, Size uint64
	Prot       uint32

	Len  uint64 `struc:"sizeof=Data"`
	Data []byte
}

type SaveBody struct {
	RegCount uint64 `struc:"sizeof=Regs"`
	Regs     []SaveReg
	MemCount uint64 `struc:"sizeof=Mem"`
	Mem      []SaveMem
}

// TODO: pack using all structs above instead of just header
func Save(u Usercorn) ([]byte, error) {
	var buf bytes.Buffer
	arch := u.Arch()
	options := &struc.Options{Order: binary.BigEndian}
	// build compressed body
	s := StrucStream{&buf, options}

	// register list
	s.Pack(uint64(len(arch.Regs)))
	for _, enum := range arch.Regs {
		val, _ := u.RegRead(enum)
		s.Pack(uint64(enum), uint64(val))
	}

	// memory mappings
	mappings := u.Mappings()
	s.Pack(uint64(len(mappings)))
	for _, m := range mappings {
		s.Pack(uint64(m.Addr), uint64(m.Size), uint32(m.Prot))
		mem, err := u.MemRead(m.Addr, m.Size)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Warning: error saving memory at 0x%x-0x%x: %s\n", m.Addr, m.Addr+m.Size, err)
			continue
		}
		buf.Write(mem)
	}

	// compress body
	var tmp bytes.Buffer
	gz := gzip.NewWriter(&tmp)
	buf.WriteTo(gz)
	buf.Reset()

	// write header / combine everything
	header := &SaveHeader{
		Version: 1,
		// unicorn version isn't exposed by Go bindings yet (Unicorn PR #483)
		UcMajor: 0, UcMinor: 0,
		UcArch: uint32(arch.UC_ARCH), UcMode: uint32(arch.UC_MODE),
		Compressed: tmp.Bytes(),
	}
	var final bytes.Buffer
	struc.PackWithOptions(&final, header, options)
	return final.Bytes(), nil
}
<|endoftext|>


<commit_before>package objectcache

import (
	"os"
	"path"
	"sort"
)

func validatePath(fileName string) bool {
	for _, char := range fileName {
		if (char >= '0' && char <= '9') || (char >= 'a' && char <= 'f') {
			continue
		}
		return false
	}
	return true
}

func cleanPath(directoryName string, fileName string) error {
	if !validatePath(fileName) {
		return os.RemoveAll(path.Join(directoryName, fileName))
	}
	return nil
}

func addCacheEntry(fileName string, cache ObjectCache) (ObjectCache, error) {
	hash, err := filenameToHash(fileName)
	if err != nil {
		return nil, err
	}
	return append(cache, hash), nil
}

func scanObjectCache(cacheDirectoryName string, subpath string,
	cache ObjectCache) (ObjectCache, error) {
	myPathName := path.Join(cacheDirectoryName, subpath)
	file, err := os.Open(myPathName)
	if err != nil {
		return nil, err
	}
	names, err := file.Readdirnames(-1)
	file.Close()
	if err != nil {
		return nil, err
	}
	for _, name := range names {
		if err = cleanPath(cacheDirectoryName, name); err != nil {
			return nil, err
		}
	}
	sort.Strings(names)
	for _, name := range names {
		fi, err := os.Lstat(path.Join(myPathName, name))
		if err != nil {
			continue
		}
		filename := path.Join(subpath, name)
		if fi.IsDir() {
			cache, err = scanObjectCache(cacheDirectoryName, filename, cache)
			if err != nil {
				return nil, err
			}
		} else {
			cache, err = addCacheEntry(filename, cache)
			if err != nil {
				return nil, err
			}
		}
	}
	return cache, nil
}
<commit_msg>Do not delete temporary/duplicate files in objectcache.<commit_after>package objectcache

import (
	"os"
	"path"
	"sort"
)

func validatePath(fileName string) bool {
	for _, char := range fileName {
		if (char >= '0' && char <= '9') || (char >= 'a' && char <= 'f') {
			continue
		}
		return false
	}
	return true
}

func addCacheEntry(fileName string, cache ObjectCache) (ObjectCache, error) {
	hash, err := filenameToHash(fileName)
	if err != nil {
		return nil, err
	}
	return append(cache, hash), nil
}

func scanObjectCache(cacheDirectoryName string, subpath string,
	cache ObjectCache) (ObjectCache, error) {
	myPathName := path.Join(cacheDirectoryName, subpath)
	file, err := os.Open(myPathName)
	if err != nil {
		return nil, err
	}
	names, err := file.Readdirnames(-1)
	file.Close()
	if err != nil {
		return nil, err
	}
	sort.Strings(names)
	for _, name := range names {
		lastChar := name[len(name)-1]
		if lastChar == '~' || lastChar == '^' {
			continue
		}
		pathname := path.Join(myPathName, name)
		if !validatePath(name) {
			if err := os.RemoveAll(pathname); err != nil {
				return nil, err
			}
			continue
		}
		fi, err := os.Lstat(pathname)
		if err != nil {
			continue
		}
		filename := path.Join(subpath, name)
		if fi.IsDir() {
			cache, err = scanObjectCache(cacheDirectoryName, filename, cache)
			if err != nil {
				return nil, err
			}
		} else {
			cache, err = addCacheEntry(filename, cache)
			if err != nil {
				return nil, err
			}
		}
	}
	return cache, nil
}
<|endoftext|>


<commit_before>// Author: Simon Labrecque <simon@wegel.ca>

package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
	"github.com/julienschmidt/httprouter"
)

var addr = flag.String("addr", ":8080", "http service address")

const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (writeWait * 5) / 10
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024 * 128,
	WriteBufferSize: 1024 * 128,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

type Hub struct {
	channels       map[uuid.UUID]*Channel
	createChannel  chan *Channel
	registerClient chan *Client
	disconnected   chan *Client
}

func newHub() *Hub {
	return &Hub{
		channels:       make(map[uuid.UUID]*Channel),
		createChannel:  make(chan *Channel),
		registerClient: make(chan *Client),
		disconnected:   make(chan *Client),
	}
}

type Channel struct {
	proxy   *Client //the proxy is on the network that we can't reach
	tunnel  *Client //the tunnel typically runs on our local computer
	id      uuid.UUID
	hub     *Hub
	handler func(*Channel)
}

func (h *Hub) setClient(client *Client) {
	if channel, ok := h.channels[client.channelID]; ok {
		if client.remoteType == "tunnel" {
			channel.tunnel = client
		} else if client.remoteType == "proxy" {
			channel.proxy = client
		}

		if channel.tunnel != nil && channel.proxy != nil {
			log.Printf("Got both sides for channel ID: %v", client.channelID.String())
			channel.tunnel.otherSide = channel.proxy
			channel.proxy.otherSide = channel.tunnel
			log.Println("Launching channel handler")
			go channel.handler(channel)
		}
	} else {
		log.Printf("Registering proxy failed for channel ID %v, channel ID unknown\n", client.channelID.String())

		go func(client *Client) {
			//tar trap potential attacker
			time.Sleep(30 * time.Second)
			client.ws.Close()
		}(client)
	}
}

func (h *Hub) handleMessages() {
	log.Println("Waiting for messages on channels")
	for {
		select {
		case channel := <-h.createChannel:
			log.Printf("Creating new channel ID: %v", channel.id.String())
			h.channels[channel.id] = channel

		//the proxy is on the network that we can't reach
		case client := <-h.registerClient:
			log.Printf("Registering %s for channel ID: %v", client.remoteType, client.channelID.String())
			h.setClient(client)

		//one of the sides disconnected, destroy the channel
		case client := <-h.disconnected:
			if channel, ok := h.channels[client.channelID]; ok {
				log.Printf("Destroying tunnel for channel ID: %v", channel.id.String())

				channel.proxy.otherSide = nil
				channel.tunnel.otherSide = nil

				channel.tunnel = nil
				channel.proxy = nil

				delete(h.channels, channel.id)
			}
		}
	}
}

func setRemote(hub *Hub, w http.ResponseWriter, r *http.Request, channelID uuid.UUID, remoteType string, params map[string][]string) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	defer ws.Close()

	client := &Client{hub: hub, ws: ws, channelID: channelID, params: params, remoteType: remoteType}
	hub.registerClient <- client
	keepalive(client)
}

func keepalive(client *Client) {
	defer func() {
		if r := recover(); r != nil {
			fmt.Printf("Exception handled in keepalive: %v\n", r)
			if client.ws != nil {
				client.ws.Close()
			}
		}
	}()

	ticker := time.NewTicker(pingPeriod)
	for {
		select {
		case <-ticker.C:
			if err := client.SetWriteDeadline(time.Now().Add(writeWait)); err != nil {
				log.Printf("error in keepalive for %s on %s: %v", client.remoteType, client.channelID, err)
			}
			if err := client.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				log.Printf("error in keepalive for %s on %s: %v", client.remoteType, client.channelID, err)
			}
		}
	}
}

func createChannel(hub *Hub, w http.ResponseWriter, r *http.Request, p httprouter.Params, channelHandler func(*Channel)) {
	log.Printf("Creating new channel")
	id := uuid.New()

	channel := &Channel{hub: hub, id: id, handler: channelHandler}
	channel.hub.createChannel <- channel

	w.Write([]byte(id.String()))
}

func serveFile(hub *Hub, w http.ResponseWriter, r *http.Request, p httprouter.Params) {
	log.Printf("Creating new channel")
	id := uuid.New()

	channel := &Channel{hub: hub, id: id}
	channel.hub.createChannel <- channel

	w.Write([]byte(id.String()))
}

func main() {
	flag.Parse()
	hub := newHub()

	router := httprouter.New()
	router.NotFound = http.FileServer(http.Dir("public"))

	router.GET("/create", func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		channelHandlerType := r.URL.Query().Get("type")
		var channelHandler func(*Channel)
		if len(channelHandlerType) == 0 || channelHandlerType == "tunnel" {
			channelHandlerType = "tunnel"
			channelHandler = Passthrough
		} else if channelHandlerType == "ssh" {
			channelHandler = sshShell
		}
		log.Println("Asked to create channel of type", channelHandlerType)
		createChannel(hub, w, r, p, channelHandler)
	})

	router.GET("/ws/proxy/:id", func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		id, _ := uuid.Parse(p.ByName("id"))
		setRemote(hub, w, r, id, "proxy", r.URL.Query())
	})
	router.GET("/ws/tunnel/:id", func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		id, _ := uuid.Parse(p.ByName("id"))
		setRemote(hub, w, r, id, "tunnel", r.URL.Query())
	})

	log.Printf("Listening on %s\n", *addr)
	log.Printf("Pinging every %v seconds\n", pingPeriod)
	go hub.handleMessages()
	log.Fatal(http.ListenAndServe(*addr, router))
}
<commit_msg>wwsconnector: use Hub.disconnected to cleanup when one side disconnects<commit_after>// Author: Simon Labrecque <simon@wegel.ca>

package main

import (
	"flag"
	"log"
	"net/http"
	"time"

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
	"github.com/julienschmidt/httprouter"
)

var addr = flag.String("addr", ":8080", "http service address")

const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (writeWait * 5) / 10
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024 * 128,
	WriteBufferSize: 1024 * 128,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

type Hub struct {
	channels       map[uuid.UUID]*Channel
	createChannel  chan *Channel
	registerClient chan *Client
	disconnected   chan *Client
}

func newHub() *Hub {
	return &Hub{
		channels:       make(map[uuid.UUID]*Channel),
		createChannel:  make(chan *Channel),
		registerClient: make(chan *Client),
		disconnected:   make(chan *Client),
	}
}

type Channel struct {
	proxy   *Client //the proxy is on the network that we can't reach
	tunnel  *Client //the tunnel typically runs on our local computer
	id      uuid.UUID
	hub     *Hub
	handler func(*Channel)
}

func (h *Hub) setClient(client *Client) {
	if channel, ok := h.channels[client.channelID]; ok {
		if client.remoteType == "tunnel" {
			channel.tunnel = client
		} else if client.remoteType == "proxy" {
			channel.proxy = client
		}

		if channel.tunnel != nil && channel.proxy != nil {
			log.Printf("Got both sides for channel ID: %v", client.channelID.String())
			channel.tunnel.otherSide = channel.proxy
			channel.proxy.otherSide = channel.tunnel
			log.Println("Launching channel handler")
			go channel.handler(channel)
		}
	} else {
		log.Printf("Registering proxy failed for channel ID %v, channel ID unknown\n", client.channelID.String())

		go func(client *Client) {
			//tar trap potential attacker
			time.Sleep(30 * time.Second)
			client.ws.Close()
		}(client)
	}
}

func (h *Hub) handleMessages() {
	log.Println("Waiting for messages on channels")
	for {
		select {
		case channel := <-h.createChannel:
			log.Printf("Creating new channel ID: %v", channel.id.String())
			h.channels[channel.id] = channel

		//the proxy is on the network that we can't reach
		case client := <-h.registerClient:
			log.Printf("Registering %s for channel ID: %v", client.remoteType, client.channelID.String())
			h.setClient(client)

		//one of the sides disconnected, destroy the channel
		case client := <-h.disconnected:
			if channel, ok := h.channels[client.channelID]; ok {
				log.Printf("Destroying tunnel for channel ID: %v", channel.id.String())
				if channel.proxy != nil {
					if channel.proxy.ws != nil {
						channel.proxy.ws.Close()
					}
					channel.proxy.otherSide = nil
					channel.proxy = nil
				}
				if channel.tunnel != nil {
					if channel.tunnel.ws != nil {
						channel.tunnel.ws.Close()
					}
					channel.tunnel.otherSide = nil
					channel.tunnel = nil
				}
				delete(h.channels, channel.id)
			}
		}
	}
}

func setRemote(hub *Hub, w http.ResponseWriter, r *http.Request, channelID uuid.UUID, remoteType string, params map[string][]string) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	defer ws.Close()

	client := &Client{hub: hub, ws: ws, channelID: channelID, params: params, remoteType: remoteType}
	hub.registerClient <- client
	keepalive(client)
}

func keepalive(client *Client) {
	defer func() {
		client.hub.disconnected <- client
		if r := recover(); r != nil {
			log.Printf("error in keepalive for %s on %s: %v", client.remoteType, client.channelID, r)
		}
	}()

	ticker := time.NewTicker(pingPeriod)
	for {
		select {
		case <-ticker.C:
			if err := client.SetWriteDeadline(time.Now().Add(writeWait)); err != nil {
				panic(err)
			}
			if err := client.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				panic(err)
			}
		}
	}
}

func createChannel(hub *Hub, w http.ResponseWriter, r *http.Request, p httprouter.Params, channelHandler func(*Channel)) {
	log.Printf("Creating new channel")
	id := uuid.New()

	channel := &Channel{hub: hub, id: id, handler: channelHandler}
	channel.hub.createChannel <- channel

	w.Write([]byte(id.String()))
}

func serveFile(hub *Hub, w http.ResponseWriter, r *http.Request, p httprouter.Params) {
	log.Printf("Creating new channel")
	id := uuid.New()

	channel := &Channel{hub: hub, id: id}
	channel.hub.createChannel <- channel

	w.Write([]byte(id.String()))
}

func main() {
	flag.Parse()
	hub := newHub()

	router := httprouter.New()
	router.NotFound = http.FileServer(http.Dir("public"))

	router.GET("/create", func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		channelHandlerType := r.URL.Query().Get("type")
		var channelHandler func(*Channel)
		if len(channelHandlerType) == 0 || channelHandlerType == "tunnel" {
			channelHandlerType = "tunnel"
			channelHandler = Passthrough
		} else if channelHandlerType == "ssh" {
			channelHandler = sshShell
		}
		log.Println("Asked to create channel of type", channelHandlerType)
		createChannel(hub, w, r, p, channelHandler)
	})

	router.GET("/ws/proxy/:id", func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		id, _ := uuid.Parse(p.ByName("id"))
		setRemote(hub, w, r, id, "proxy", r.URL.Query())
	})
	router.GET("/ws/tunnel/:id", func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		id, _ := uuid.Parse(p.ByName("id"))
		setRemote(hub, w, r, id, "tunnel", r.URL.Query())
	})

	log.Printf("Listening on %s\n", *addr)
	log.Printf("Pinging every %v seconds\n", pingPeriod)
	go hub.handleMessages()
	log.Fatal(http.ListenAndServe(*addr, router))
}
<|endoftext|>


<commit_before>package pnet

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"io"
)

func newLine() io.Reader {
	return bytes.NewReader([]byte("\n"))
}

func GenerateV1PSK() io.Reader {
	psk := make([]byte, 32)
	rand.Read(psk)
	hexPsk := make([]byte, len(psk)*2)
	hex.Encode(hexPsk, psk)

	// just a shortcut to NewReader
	nr := func(b []byte) io.Reader {
		return bytes.NewReader(b)
	}
	return io.MultiReader(nr(pathPSKv1), newLine(), nr([]byte("/base16/")), newLine(), nr(hexPsk))
}
<commit_msg>Add docs to GenerateV1PSK<commit_after>package pnet

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"io"
)

func newLine() io.Reader {
	return bytes.NewReader([]byte("\n"))
}

// GenerateV1PSK generates new PSK key that can be used with NewProtector
func GenerateV1PSK() io.Reader {
	psk := make([]byte, 32)
	rand.Read(psk)
	hexPsk := make([]byte, len(psk)*2)
	hex.Encode(hexPsk, psk)

	// just a shortcut to NewReader
	nr := func(b []byte) io.Reader {
		return bytes.NewReader(b)
	}
	return io.MultiReader(nr(pathPSKv1), newLine(), nr([]byte("/base16/")), newLine(), nr(hexPsk))
}
<|endoftext|>


<commit_before>// Copyright 2015 The rkt 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.

//+build linux

package main

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"

	"github.com/coreos/rkt/Godeps/_workspace/src/github.com/appc/spec/schema/types"
	"github.com/coreos/rkt/Godeps/_workspace/src/github.com/coreos/go-systemd/unit"
)

type addIsolatorFunc func(opts []*unit.UnitOption, limit string) ([]*unit.UnitOption, error)

var (
	isolatorFuncs = map[string]addIsolatorFunc{
		"cpu":    addCpuLimit,
		"memory": addMemoryLimit,
	}
	cgroupControllerRWFiles = map[string][]string{
		"memory": []string{"memory.limit_in_bytes"},
		"cpu":    []string{"cpu.cfs_quota_us"},
	}
)

func addCpuLimit(opts []*unit.UnitOption, limit string) ([]*unit.UnitOption, error) {
	milliCores, err := strconv.Atoi(limit)
	if err != nil {
		return nil, err
	}
	quota := strconv.Itoa(milliCores/10) + "%"
	opts = append(opts, newUnitOption("Service", "CPUQuota", quota))
	return opts, nil
}

func addMemoryLimit(opts []*unit.UnitOption, limit string) ([]*unit.UnitOption, error) {
	opts = append(opts, newUnitOption("Service", "MemoryLimit", limit))
	return opts, nil
}

func maybeAddIsolator(opts []*unit.UnitOption, isolator string, limit string) ([]*unit.UnitOption, error) {
	var err error
	if isIsolatorSupported(isolator) {
		opts, err = isolatorFuncs[isolator](opts, limit)
		if err != nil {
			return nil, err
		}
	} else {
		fmt.Fprintf(os.Stderr, "warning: resource/%s isolator set but support disabled in the kernel, skipping\n", isolator)
	}
	return opts, nil
}

func isIsolatorSupported(isolator string) bool {
	if files, ok := cgroupControllerRWFiles[isolator]; ok {
		for _, f := range files {
			isolatorPath := filepath.Join("/sys/fs/cgroup/", isolator, f)
			if _, err := os.Stat(isolatorPath); os.IsNotExist(err) {
				return false
			}
		}
		return true
	}
	return false
}

func parseCgroups(f io.Reader) (map[int][]string, error) {
	sc := bufio.NewScanner(f)

	// skip first line since it is a comment
	sc.Scan()

	cgroups := make(map[int][]string)
	for sc.Scan() {
		var controller string
		var hierarchy int
		var num int
		var enabled int
		fmt.Sscanf(sc.Text(), "%s %d %d %d", &controller, &hierarchy, &num, &enabled)

		if enabled == 1 {
			if _, ok := cgroups[hierarchy]; !ok {
				cgroups[hierarchy] = []string{controller}
			} else {
				cgroups[hierarchy] = append(cgroups[hierarchy], controller)
			}
		}
	}

	if err := sc.Err(); err != nil {
		return nil, err
	}

	return cgroups, nil
}

func getControllers(cgroups map[int][]string) []string {
	var controllers []string
	for _, cs := range cgroups {
		controllers = append(controllers, strings.Join(cs, ","))
	}

	return controllers
}

func getControllerSymlinks(cgroups map[int][]string) map[string]string {
	symlinks := make(map[string]string)

	for _, cs := range cgroups {
		if len(cs) > 1 {
			tgt := strings.Join(cs, ",")
			for _, ln := range cs {
				symlinks[ln] = tgt
			}
		}
	}

	return symlinks
}

func getControllerRWFiles(controller string) []string {
	parts := strings.Split(controller, ",")
	for _, p := range parts {
		if files, ok := cgroupControllerRWFiles[p]; ok {
			// cgroup.procs always needs to be RW for allowing systemd to add
			// processes to the controller
			files = append(files, "cgroup.procs")
			return files
		}
	}

	return nil
}

func getOwnCgroupPath(controller string) (string, error) {
	selfCgroupPath := "/proc/self/cgroup"
	cg, err := os.Open(selfCgroupPath)
	if err != nil {
		return "", fmt.Errorf("error opening /proc/self/cgroup: %v", err)
	}
	defer cg.Close()

	s := bufio.NewScanner(cg)
	for s.Scan() {
		parts := strings.Split(s.Text(), ":")
		if parts[1] == controller {
			return parts[2], nil
		}
	}

	return "", fmt.Errorf("controller %q not found", controller)
}

// createCgroups mounts the cgroup controllers hierarchy for the container but
// leaves the subcgroup for each app read-write so the systemd inside stage1
// can apply isolators to them
func createCgroups(root string, subcgroup string, appHashes []types.Hash) error {
	cgroupsFile, err := os.Open("/proc/cgroups")
	if err != nil {
		return err
	}
	defer cgroupsFile.Close()

	cgroups, err := parseCgroups(cgroupsFile)
	if err != nil {
		return fmt.Errorf("error parsing /proc/cgroups: %v", err)
	}

	controllers := getControllers(cgroups)

	var flags uintptr

	// 1. Mount /sys read-only
	sys := filepath.Join(root, "/sys")
	if err := os.MkdirAll(sys, 0700); err != nil {
		return err
	}
	flags = syscall.MS_RDONLY |
		syscall.MS_NOSUID |
		syscall.MS_NOEXEC |
		syscall.MS_NODEV
	if err := syscall.Mount("sysfs", sys, "sysfs", flags, ""); err != nil {
		return fmt.Errorf("error mounting %q: %v", sys, err)
	}

	// 2. Mount /sys/fs/cgroup
	cgroupTmpfs := filepath.Join(root, "/sys/fs/cgroup")
	if err := os.MkdirAll(cgroupTmpfs, 0700); err != nil {
		return err
	}

	flags = syscall.MS_NOSUID |
		syscall.MS_NOEXEC |
		syscall.MS_NODEV |
		syscall.MS_STRICTATIME
	if err := syscall.Mount("tmpfs", cgroupTmpfs, "tmpfs", flags, "mode=755"); err != nil {
		return fmt.Errorf("error mounting %q: %v", cgroupTmpfs, err)
	}

	// 3. Mount controllers
	for _, c := range controllers {
		// 3a. Mount controller
		cPath := filepath.Join(root, "/sys/fs/cgroup", c)
		if err := os.MkdirAll(cPath, 0700); err != nil {
			return err
		}

		flags = syscall.MS_NOSUID |
			syscall.MS_NOEXEC |
			syscall.MS_NODEV
		if err := syscall.Mount("cgroup", cPath, "cgroup", flags, c); err != nil {
			return fmt.Errorf("error mounting %q: %v", cPath, err)
		}

		// 3b. Check if we're running from a unit to know which subcgroup
		// directories to mount read-write
		subcgroupPath := filepath.Join(cPath, subcgroup)

		// 3c. Create cgroup directories and mount the files we need over
		// themselves so they stay read-write
		for _, a := range appHashes {
			serviceName := ServiceUnitName(a)
			appCgroup := filepath.Join(subcgroupPath, serviceName)
			if err := os.MkdirAll(appCgroup, 0755); err != nil {
				return err
			}
			for _, f := range getControllerRWFiles(c) {
				cgroupFilePath := filepath.Join(appCgroup, f)
				// the file may not be there if kernel doesn't support the
				// feature, skip it in that case
				if _, err := os.Stat(cgroupFilePath); os.IsNotExist(err) {
					continue
				}
				if err := syscall.Mount(cgroupFilePath, cgroupFilePath, "", syscall.MS_BIND, ""); err != nil {
					return fmt.Errorf("error bind mounting %q: %v", cgroupFilePath, err)
				}
			}
		}

		// 3d. Re-mount controller read-only to prevent the container modifying host controllers
		flags = syscall.MS_BIND |
			syscall.MS_REMOUNT |
			syscall.MS_NOSUID |
			syscall.MS_NOEXEC |
			syscall.MS_NODEV |
			syscall.MS_RDONLY
		if err := syscall.Mount(cPath, cPath, "", flags, ""); err != nil {
			return fmt.Errorf("error remounting RO %q: %v", cPath, err)
		}
	}

	// 4. Create symlinks for combined controllers
	symlinks := getControllerSymlinks(cgroups)
	for ln, tgt := range symlinks {
		lnPath := filepath.Join(cgroupTmpfs, ln)
		if err := os.Symlink(tgt, lnPath); err != nil {
			return fmt.Errorf("error creating symlink: %v", err)
		}
	}

	// 5. Create systemd cgroup directory
	// We're letting systemd-nspawn create the systemd cgroup but later we're
	// remounting /sys/fs/cgroup read-only so we create the directory here.
	if err := os.MkdirAll(filepath.Join(cgroupTmpfs, "systemd"), 0700); err != nil {
		return err
	}

	// 6. Bind-mount cgroup filesystem read-only
	flags = syscall.MS_BIND |
		syscall.MS_REMOUNT |
		syscall.MS_NOSUID |
		syscall.MS_NOEXEC |
		syscall.MS_NODEV |
		syscall.MS_RDONLY
	if err := syscall.Mount(cgroupTmpfs, cgroupTmpfs, "", flags, ""); err != nil {
		return fmt.Errorf("error remounting RO %q: %v", cgroupTmpfs, err)
	}

	return nil
}
<commit_msg>stage1: make getOwnCgroupPath() more robust<commit_after>// Copyright 2015 The rkt 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.

//+build linux

package main

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"

	"github.com/coreos/rkt/Godeps/_workspace/src/github.com/appc/spec/schema/types"
	"github.com/coreos/rkt/Godeps/_workspace/src/github.com/coreos/go-systemd/unit"
)

type addIsolatorFunc func(opts []*unit.UnitOption, limit string) ([]*unit.UnitOption, error)

var (
	isolatorFuncs = map[string]addIsolatorFunc{
		"cpu":    addCpuLimit,
		"memory": addMemoryLimit,
	}
	cgroupControllerRWFiles = map[string][]string{
		"memory": []string{"memory.limit_in_bytes"},
		"cpu":    []string{"cpu.cfs_quota_us"},
	}
)

func addCpuLimit(opts []*unit.UnitOption, limit string) ([]*unit.UnitOption, error) {
	milliCores, err := strconv.Atoi(limit)
	if err != nil {
		return nil, err
	}
	quota := strconv.Itoa(milliCores/10) + "%"
	opts = append(opts, newUnitOption("Service", "CPUQuota", quota))
	return opts, nil
}

func addMemoryLimit(opts []*unit.UnitOption, limit string) ([]*unit.UnitOption, error) {
	opts = append(opts, newUnitOption("Service", "MemoryLimit", limit))
	return opts, nil
}

func maybeAddIsolator(opts []*unit.UnitOption, isolator string, limit string) ([]*unit.UnitOption, error) {
	var err error
	if isIsolatorSupported(isolator) {
		opts, err = isolatorFuncs[isolator](opts, limit)
		if err != nil {
			return nil, err
		}
	} else {
		fmt.Fprintf(os.Stderr, "warning: resource/%s isolator set but support disabled in the kernel, skipping\n", isolator)
	}
	return opts, nil
}

func isIsolatorSupported(isolator string) bool {
	if files, ok := cgroupControllerRWFiles[isolator]; ok {
		for _, f := range files {
			isolatorPath := filepath.Join("/sys/fs/cgroup/", isolator, f)
			if _, err := os.Stat(isolatorPath); os.IsNotExist(err) {
				return false
			}
		}
		return true
	}
	return false
}

func parseCgroups(f io.Reader) (map[int][]string, error) {
	sc := bufio.NewScanner(f)

	// skip first line since it is a comment
	sc.Scan()

	cgroups := make(map[int][]string)
	for sc.Scan() {
		var controller string
		var hierarchy int
		var num int
		var enabled int
		fmt.Sscanf(sc.Text(), "%s %d %d %d", &controller, &hierarchy, &num, &enabled)

		if enabled == 1 {
			if _, ok := cgroups[hierarchy]; !ok {
				cgroups[hierarchy] = []string{controller}
			} else {
				cgroups[hierarchy] = append(cgroups[hierarchy], controller)
			}
		}
	}

	if err := sc.Err(); err != nil {
		return nil, err
	}

	return cgroups, nil
}

func getControllers(cgroups map[int][]string) []string {
	var controllers []string
	for _, cs := range cgroups {
		controllers = append(controllers, strings.Join(cs, ","))
	}

	return controllers
}

func getControllerSymlinks(cgroups map[int][]string) map[string]string {
	symlinks := make(map[string]string)

	for _, cs := range cgroups {
		if len(cs) > 1 {
			tgt := strings.Join(cs, ",")
			for _, ln := range cs {
				symlinks[ln] = tgt
			}
		}
	}

	return symlinks
}

func getControllerRWFiles(controller string) []string {
	parts := strings.Split(controller, ",")
	for _, p := range parts {
		if files, ok := cgroupControllerRWFiles[p]; ok {
			// cgroup.procs always needs to be RW for allowing systemd to add
			// processes to the controller
			files = append(files, "cgroup.procs")
			return files
		}
	}

	return nil
}

func getOwnCgroupPath(controller string) (string, error) {
	selfCgroupPath := "/proc/self/cgroup"
	cg, err := os.Open(selfCgroupPath)
	if err != nil {
		return "", fmt.Errorf("error opening /proc/self/cgroup: %v", err)
	}
	defer cg.Close()

	s := bufio.NewScanner(cg)
	for s.Scan() {
		parts := strings.SplitN(s.Text(), ":", 3)
		if len(parts) < 3 {
			return "", fmt.Errorf("error parsing /proc/self/cgroup")
		}
		if parts[1] == controller {
			return parts[2], nil
		}
	}

	return "", fmt.Errorf("controller %q not found", controller)
}

// createCgroups mounts the cgroup controllers hierarchy for the container but
// leaves the subcgroup for each app read-write so the systemd inside stage1
// can apply isolators to them
func createCgroups(root string, subcgroup string, appHashes []types.Hash) error {
	cgroupsFile, err := os.Open("/proc/cgroups")
	if err != nil {
		return err
	}
	defer cgroupsFile.Close()

	cgroups, err := parseCgroups(cgroupsFile)
	if err != nil {
		return fmt.Errorf("error parsing /proc/cgroups: %v", err)
	}

	controllers := getControllers(cgroups)

	var flags uintptr

	// 1. Mount /sys read-only
	sys := filepath.Join(root, "/sys")
	if err := os.MkdirAll(sys, 0700); err != nil {
		return err
	}
	flags = syscall.MS_RDONLY |
		syscall.MS_NOSUID |
		syscall.MS_NOEXEC |
		syscall.MS_NODEV
	if err := syscall.Mount("sysfs", sys, "sysfs", flags, ""); err != nil {
		return fmt.Errorf("error mounting %q: %v", sys, err)
	}

	// 2. Mount /sys/fs/cgroup
	cgroupTmpfs := filepath.Join(root, "/sys/fs/cgroup")
	if err := os.MkdirAll(cgroupTmpfs, 0700); err != nil {
		return err
	}

	flags = syscall.MS_NOSUID |
		syscall.MS_NOEXEC |
		syscall.MS_NODEV |
		syscall.MS_STRICTATIME
	if err := syscall.Mount("tmpfs", cgroupTmpfs, "tmpfs", flags, "mode=755"); err != nil {
		return fmt.Errorf("error mounting %q: %v", cgroupTmpfs, err)
	}

	// 3. Mount controllers
	for _, c := range controllers {
		// 3a. Mount controller
		cPath := filepath.Join(root, "/sys/fs/cgroup", c)
		if err := os.MkdirAll(cPath, 0700); err != nil {
			return err
		}

		flags = syscall.MS_NOSUID |
			syscall.MS_NOEXEC |
			syscall.MS_NODEV
		if err := syscall.Mount("cgroup", cPath, "cgroup", flags, c); err != nil {
			return fmt.Errorf("error mounting %q: %v", cPath, err)
		}

		// 3b. Check if we're running from a unit to know which subcgroup
		// directories to mount read-write
		subcgroupPath := filepath.Join(cPath, subcgroup)

		// 3c. Create cgroup directories and mount the files we need over
		// themselves so they stay read-write
		for _, a := range appHashes {
			serviceName := ServiceUnitName(a)
			appCgroup := filepath.Join(subcgroupPath, serviceName)
			if err := os.MkdirAll(appCgroup, 0755); err != nil {
				return err
			}
			for _, f := range getControllerRWFiles(c) {
				cgroupFilePath := filepath.Join(appCgroup, f)
				// the file may not be there if kernel doesn't support the
				// feature, skip it in that case
				if _, err := os.Stat(cgroupFilePath); os.IsNotExist(err) {
					continue
				}
				if err := syscall.Mount(cgroupFilePath, cgroupFilePath, "", syscall.MS_BIND, ""); err != nil {
					return fmt.Errorf("error bind mounting %q: %v", cgroupFilePath, err)
				}
			}
		}

		// 3d. Re-mount controller read-only to prevent the container modifying host controllers
		flags = syscall.MS_BIND |
			syscall.MS_REMOUNT |
			syscall.MS_NOSUID |
			syscall.MS_NOEXEC |
			syscall.MS_NODEV |
			syscall.MS_RDONLY
		if err := syscall.Mount(cPath, cPath, "", flags, ""); err != nil {
			return fmt.Errorf("error remounting RO %q: %v", cPath, err)
		}
	}

	// 4. Create symlinks for combined controllers
	symlinks := getControllerSymlinks(cgroups)
	for ln, tgt := range symlinks {
		lnPath := filepath.Join(cgroupTmpfs, ln)
		if err := os.Symlink(tgt, lnPath); err != nil {
			return fmt.Errorf("error creating symlink: %v", err)
		}
	}

	// 5. Create systemd cgroup directory
	// We're letting systemd-nspawn create the systemd cgroup but later we're
	// remounting /sys/fs/cgroup read-only so we create the directory here.
	if err := os.MkdirAll(filepath.Join(cgroupTmpfs, "systemd"), 0700); err != nil {
		return err
	}

	// 6. Bind-mount cgroup filesystem read-only
	flags = syscall.MS_BIND |
		syscall.MS_REMOUNT |
		syscall.MS_NOSUID |
		syscall.MS_NOEXEC |
		syscall.MS_NODEV |
		syscall.MS_RDONLY
	if err := syscall.Mount(cgroupTmpfs, cgroupTmpfs, "", flags, ""); err != nil {
		return fmt.Errorf("error remounting RO %q: %v", cgroupTmpfs, err)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"testing"
)

func TestTrivial(t *testing.T) {
	got := 1 + 1
	want := 2
	if got != want {
		t.Errorf("1 + 1 is %v, want %v", got, want)
	}
}
<commit_msg>Delete a dummy test<commit_after><|endoftext|>