<commit_before>package main

import (
	"html/template"

	"github.com/youtube/vitess/go/vt/health"
	"github.com/youtube/vitess/go/vt/servenv"
	_ "github.com/youtube/vitess/go/vt/status"
	"github.com/youtube/vitess/go/vt/tabletmanager"
	"github.com/youtube/vitess/go/vt/tabletserver"
	"github.com/youtube/vitess/go/vt/topo"
)

var (
	// tabletTemplate contains the style sheet and the tablet itself.
	tabletTemplate = `

      Alias: {{github_com_youtube_vitess_vtctld_tablet .Tablet.AliasString}}

      Keyspace: {{github_com_youtube_vitess_vtctld_keyspace .Tablet.Keyspace}} Shard: {{github_com_youtube_vitess_vtctld_shard .Tablet.Keyspace .Tablet.Shard}} Tablet Type: {{.Tablet.Type}}

      SrvKeyspace: {{github_com_youtube_vitess_vtctld_srv_keyspace .Tablet.Alias.Cell .Tablet.Keyspace}}

      Replication graph: {{github_com_youtube_vitess_vtctld_replication .Tablet.Alias.Cell .Tablet.Keyspace .Tablet.Shard}}

      {{if .BlacklistedTables}}
        BlacklistedTables: {{range .BlacklistedTables}}{{.}} {{end}}

      {{end}}
      {{if .DisallowQueryService}}
        Query Service disabled: {{.DisallowQueryService}}

      {{end}}
      {{if .DisableUpdateStream}}
        Update Stream disabled

      {{end}}

      Schema

      Schema Query Plans

      Schema Query Stats

      Schema Table Stats

      Query Stats

      Consolidations

      Current Query Log

      Current Transaction Log

      Health Check

      Query Service Health Check

      Current Stream Queries

`

	// healthTemplate is just about the tablet health
	healthTemplate = `

Current status: {{.CurrentHTML}}

Polling health information from {{github_com_youtube_vitess_health_html_name}}. ({{.Config}})

Health History

    Time
    Healthcheck Result

  {{range .Records}}

    {{.Time.Format "Jan 2, 2006 at 15:04:05 (MST)"}}
    {{.HTML}}

  {{end}}

healthy

serving traffic.

unhappy

will serve traffic only if there are no fully healthy tablets.

unhealthy

will not serve traffic.

`

	// binlogTemplate is about the binlog players
	binlogTemplate = `
{{if .Controllers}}
Binlog player state: {{.State}}

    Index
    SourceShard
    State
    StopPosition
    LastPosition
    SecondsBehindMaster
    Counts
    Rates
    Last Error

  {{range .Controllers}}

      {{.Index}}
      {{.SourceShardAsHTML}}
      {{.State}}
        {{if eq .State "Running"}}
          {{if .SourceTabletAlias}}
            (from {{github_com_youtube_vitess_vtctld_tablet .SourceTabletAlias}})
          {{else}}
            (picking source tablet)
          {{end}}
        {{end}}
      {{if .StopPosition}}{{.StopPosition}}{{end}}
      {{.LastPosition}}
      {{.SecondsBehindMaster}}
      {{range $key, $value := .Counts}}{{$key}}: {{$value}}
{{end}}
      {{range $key, $values := .Rates}}{{$key}}: {{range $values}}{{.}} {{end}}
{{end}}
      {{.LastError}}

  {{end}}

{{else}}
No binlog player is running.
{{end}}
`
)

type healthStatus struct {
	Records []interface{}
	Config  template.HTML
}

func (hs *healthStatus) CurrentClass() string {
	if len(hs.Records) > 0 {
		return hs.Records[0].(*tabletmanager.HealthRecord).Class()
	}
	return "unknown"
}

func (hs *healthStatus) CurrentHTML() template.HTML {
	if len(hs.Records) > 0 {
		return hs.Records[0].(*tabletmanager.HealthRecord).HTML()
	}
	return template.HTML("unknown")
}

func healthHTMLName() template.HTML {
	return health.DefaultAggregator.HTMLName()
}

// For use by plugins which wish to avoid racing when registering status page parts.
var onStatusRegistered func()

func addStatusParts(qsc tabletserver.Controller) {
	servenv.AddStatusPart("Tablet", tabletTemplate, func() interface{} {
		return map[string]interface{}{
			"Tablet":               topo.NewTabletInfo(agent.Tablet(), -1),
			"BlacklistedTables":    agent.BlacklistedTables(),
			"DisallowQueryService": agent.DisallowQueryService(),
			"DisableUpdateStream":  !agent.EnableUpdateStream(),
		}
	})
	servenv.AddStatusFuncs(template.FuncMap{
		"github_com_youtube_vitess_health_html_name": healthHTMLName,
	})
	servenv.AddStatusPart("Health", healthTemplate, func() interface{} {
		return &healthStatus{
			Records: agent.History.Records(),
			Config:  tabletmanager.ConfigHTML(),
		}
	})
	qsc.AddStatusPart()
	servenv.AddStatusPart("Binlog Player", binlogTemplate, func() interface{} {
		return agent.BinlogPlayerMap.Status()
	})
	if onStatusRegistered != nil {
		onStatusRegistered()
	}
}
<commit_msg>b/30066332: export /streamqueryz in status page<commit_after>package main

import (
	"html/template"

	"github.com/youtube/vitess/go/vt/health"
	"github.com/youtube/vitess/go/vt/servenv"
	_ "github.com/youtube/vitess/go/vt/status"
	"github.com/youtube/vitess/go/vt/tabletmanager"
	"github.com/youtube/vitess/go/vt/tabletserver"
	"github.com/youtube/vitess/go/vt/topo"
)

var (
	// tabletTemplate contains the style sheet and the tablet itself.
	tabletTemplate = `

      Alias: {{github_com_youtube_vitess_vtctld_tablet .Tablet.AliasString}}

      Keyspace: {{github_com_youtube_vitess_vtctld_keyspace .Tablet.Keyspace}} Shard: {{github_com_youtube_vitess_vtctld_shard .Tablet.Keyspace .Tablet.Shard}} Tablet Type: {{.Tablet.Type}}

      SrvKeyspace: {{github_com_youtube_vitess_vtctld_srv_keyspace .Tablet.Alias.Cell .Tablet.Keyspace}}

      Replication graph: {{github_com_youtube_vitess_vtctld_replication .Tablet.Alias.Cell .Tablet.Keyspace .Tablet.Shard}}

      {{if .BlacklistedTables}}
        BlacklistedTables: {{range .BlacklistedTables}}{{.}} {{end}}

      {{end}}
      {{if .DisallowQueryService}}
        Query Service disabled: {{.DisallowQueryService}}

      {{end}}
      {{if .DisableUpdateStream}}
        Update Stream disabled

      {{end}}

      Schema

      Schema Query Plans

      Schema Query Stats

      Schema Table Stats

      Query Stats

      Streaming Query Stats

      Consolidations

      Current Query Log

      Current Transaction Log

      Health Check

      Query Service Health Check

      Current Stream Queries

`

	// healthTemplate is just about the tablet health
	healthTemplate = `

Current status: {{.CurrentHTML}}

Polling health information from {{github_com_youtube_vitess_health_html_name}}. ({{.Config}})

Health History

    Time
    Healthcheck Result

  {{range .Records}}

    {{.Time.Format "Jan 2, 2006 at 15:04:05 (MST)"}}
    {{.HTML}}

  {{end}}

healthy

serving traffic.

unhappy

will serve traffic only if there are no fully healthy tablets.

unhealthy

will not serve traffic.

`

	// binlogTemplate is about the binlog players
	binlogTemplate = `
{{if .Controllers}}
Binlog player state: {{.State}}

    Index
    SourceShard
    State
    StopPosition
    LastPosition
    SecondsBehindMaster
    Counts
    Rates
    Last Error

  {{range .Controllers}}

      {{.Index}}
      {{.SourceShardAsHTML}}
      {{.State}}
        {{if eq .State "Running"}}
          {{if .SourceTabletAlias}}
            (from {{github_com_youtube_vitess_vtctld_tablet .SourceTabletAlias}})
          {{else}}
            (picking source tablet)
          {{end}}
        {{end}}
      {{if .StopPosition}}{{.StopPosition}}{{end}}
      {{.LastPosition}}
      {{.SecondsBehindMaster}}
      {{range $key, $value := .Counts}}{{$key}}: {{$value}}
{{end}}
      {{range $key, $values := .Rates}}{{$key}}: {{range $values}}{{.}} {{end}}
{{end}}
      {{.LastError}}

  {{end}}

{{else}}
No binlog player is running.
{{end}}
`
)

type healthStatus struct {
	Records []interface{}
	Config  template.HTML
}

func (hs *healthStatus) CurrentClass() string {
	if len(hs.Records) > 0 {
		return hs.Records[0].(*tabletmanager.HealthRecord).Class()
	}
	return "unknown"
}

func (hs *healthStatus) CurrentHTML() template.HTML {
	if len(hs.Records) > 0 {
		return hs.Records[0].(*tabletmanager.HealthRecord).HTML()
	}
	return template.HTML("unknown")
}

func healthHTMLName() template.HTML {
	return health.DefaultAggregator.HTMLName()
}

// For use by plugins which wish to avoid racing when registering status page parts.
var onStatusRegistered func()

func addStatusParts(qsc tabletserver.Controller) {
	servenv.AddStatusPart("Tablet", tabletTemplate, func() interface{} {
		return map[string]interface{}{
			"Tablet":               topo.NewTabletInfo(agent.Tablet(), -1),
			"BlacklistedTables":    agent.BlacklistedTables(),
			"DisallowQueryService": agent.DisallowQueryService(),
			"DisableUpdateStream":  !agent.EnableUpdateStream(),
		}
	})
	servenv.AddStatusFuncs(template.FuncMap{
		"github_com_youtube_vitess_health_html_name": healthHTMLName,
	})
	servenv.AddStatusPart("Health", healthTemplate, func() interface{} {
		return &healthStatus{
			Records: agent.History.Records(),
			Config:  tabletmanager.ConfigHTML(),
		}
	})
	qsc.AddStatusPart()
	servenv.AddStatusPart("Binlog Player", binlogTemplate, func() interface{} {
		return agent.BinlogPlayerMap.Status()
	})
	if onStatusRegistered != nil {
		onStatusRegistered()
	}
}
<|endoftext|>


<commit_before>/*
   Copyright 2014 Outbrain 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

       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 discovery manages a queue of discovery requests: an ordered
queue with no duplicates.

The configured processor function is called on an entry that's
received and this happens in parallel until maxCapacity go-routines
are running in parallel.

Any new requests get queued and processed as a go routine becomes
free.  The queue is processed in FIFO order.  If a request is
received for an instance that is already in the queue or already
being processed then it will be silently ignored.

*/

package discovery

import (
	"errors"
	"sync"

	"github.com/outbrain/golib/log"
	"github.com/outbrain/orchestrator/go/inst"
)

// Queue is a container for processing the orchestrator discovery requests.
type Queue struct {
	// The number of active go routines processing discovery requests.
	concurrency uint
	// Channel used by the active go routines to say they have finished processing.
	done chan inst.InstanceKey
	// Input channel we are reading the discovery requests from.
	inputChan <-chan inst.InstanceKey
	// This has 2 uses: to indicate if there is a request in the queue
	// (so not being processed) or to indicate that the request is actively
	// being dealth with.  That state is not explicitly stored as it is not
	// really needed.
	knownKeys map[inst.InstanceKey]bool
	// lock while making changes
	lock sync.Mutex
	// The maximum number of go routines allowed to handle the queue at once.
	// This is a configuration parameter provided when creating the Queue.
	maxConcurrency uint
	// process to run on each received key
	processor func(i inst.InstanceKey)
	// This holds the discover requests (the instance name) which need to be
	// processed. but which are not currently being processed. All requests
	// are initially added to the end of this queue, and then the top element
	// will be popped off if the number of active go routines (defined by
	// concurrency) is less than the maximum specified value at which point
	// it will be processed by a new go routine.
	queue []inst.InstanceKey
}

var emptyKey = inst.InstanceKey{}

// NewQueue creates a new Queue entry.
func NewQueue(maxConcurrency uint, inputChan chan inst.InstanceKey, processor func(i inst.InstanceKey)) *Queue {
	log.Infof("Queue.NewQueue()")
	q := new(Queue)

	q.concurrency = 0                    // explicitly
	q.done = make(chan inst.InstanceKey) // Do I need this to be larger?
	q.inputChan = inputChan
	q.knownKeys = make(map[inst.InstanceKey]bool)
	q.maxConcurrency = maxConcurrency
	q.processor = processor
	q.queue = make([]inst.InstanceKey, 0)

	return q
}

// add the key to the slice if it does not exist in known keys
// - goroutine safe as only called inside the mutex
func (q *Queue) push(key inst.InstanceKey) {
	if key == emptyKey {
		log.Fatal("Queue.push(%v) is empty", key)
	}
	// log.Debugf("Queue.push(%+v)", key)

	if _, found := q.knownKeys[key]; !found {
		// log.Debugf("Queue.push() adding %+v to knownKeys", key)
		// add to the items that are being processed
		q.knownKeys[key] = true
		q.queue = append(q.queue, key)
	} else {
		// If key already there we just ignore it as the request is in the queue.
		// the known key also records stuff in the queue, so pending + active jobs.
		// log.Debugf("Queue.push() ignoring knownKey %+v", key)
	}
}

// remove the entry and remove it from known keys
func (q *Queue) pop() (inst.InstanceKey, error) {
	if len(q.queue) == 0 {
		return inst.InstanceKey{}, errors.New("q.pop() on empty queue")
	}
	key := q.queue[0]
	q.queue = q.queue[1:]
	delete(q.knownKeys, key)
	// log.Debugf("Queue.pop() returns %+v", key)
	return key, nil
}

// dispatch a job from the queue (assumes we are in a locked state)
func (q *Queue) dispatch() {
	key, err := q.pop() // should never give an error but let's check anyway
	if err != nil {
		log.Fatal("Queue.dispatch() q.pop() returns: %+v", err)
		return
	}
	if key == emptyKey {
		log.Fatal("Queue.dispatch() key is empty")
	}

	q.concurrency++
	q.knownKeys[key] = true

	// log.Debugf("Queue.dispatch() key: %q, concurrency: %d", key, q.concurrency)

	// dispatch a discoverInstance() but tell us when we're done (to limit concurrency)
	go func() { // discover asynchronously
		q.processor(key)
		q.done <- key
	}()
}

// acknowledge a job has finished
// - we deal with the locking inside
func (q *Queue) acknowledgeJob(key inst.InstanceKey) {
	q.lock.Lock()
	delete(q.knownKeys, key)
	q.concurrency--
	// log.Debugf("Queue.acknowledgeJob(%+v) q.concurrency: %d", key, q.concurrency)
	q.lock.Unlock()
}

// drain queue by dispatching any jobs we have still
func (q *Queue) maybeDispatch() {
	q.lock.Lock()
	// log.Debugf("Queue.maybeDispatch() q.concurrency: %d, q.maxConcurrency: %d, len(q.queue): %d", q.concurrency, q.maxConcurrency, len(q.queue))
	if q.concurrency < q.maxConcurrency && len(q.queue) > 0 {
		q.dispatch()
	}
	q.lock.Unlock()
}

// add an entry to the queue and dispatch something if concurrency is low enough
// - we deal with locking inside
func (q *Queue) queueAndMaybeDispatch(key inst.InstanceKey) {
	if key == emptyKey {
		log.Fatal("Queue.queueAndMaybeDispatch(%v) is empty", key)
	}
	q.lock.Lock()
	// log.Debugf("Queue.queueAndMaybeDispatch(%+v) concurency: %d", key, q.concurrency)
	q.push(key)
	if q.concurrency < q.maxConcurrency && len(q.queue) > 0 {
		q.dispatch()
	}
	q.lock.Unlock()
}

// cleanup is called when the input channel closes.
// we can not sit in the loop so we have to wait for running go-routines to finish
// but also to dispatch anything left in the queue until finally everything is done.
func (q *Queue) cleanup() {
	log.Infof("Queue.cleanup()")
	for q.concurrency > 0 && len(q.queue) > 0 {
		q.maybeDispatch()
		if key, closed := <-q.done; closed {
			return
		} else {
			q.acknowledgeJob(key)
		}
	}
}

// Ends when all elements in the queue have been handled.
// we read from inputChan and call processor up to maxConcurrency times in parallel
func (q *Queue) HandleRequests() {
	if q == nil {
		log.Infof("Queue.HandleRequests() q == nil ??. Should not happen")

		// no queue, nothing to do
		return
	}
	log.Infof("Queue.NewQueue() processing requests")
	for {
		select {
		case key, ok := <-q.inputChan:
			if ok {
				if key != emptyKey {
					q.queueAndMaybeDispatch(key)
				} else {
					log.Warningf("Queue.HandleRequests() q.inputChan received empty key %+v, ignoring (fix the upstream code to prevent this)", key)
				}
			} else {
				q.cleanup()
				log.Infof("Queue.HandleRequests() q.inputChan is closed. returning")
				return
			}
		case key, ok := <-q.done:
			if ok {
				q.acknowledgeJob(key)
			} else {
				log.Infof("Queue.HandleRequests() q.done is closed. returning (shouldn't get here)")
				return // we shouldn't get here as the return above should get triggered first
			}
		}
	}
}
<commit_msg>rm debug comments<commit_after>/*
   Copyright 2014 Outbrain 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

       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 discovery manages a queue of discovery requests: an ordered
queue with no duplicates.

The configured processor function is called on an entry that's
received and this happens in parallel until maxCapacity go-routines
are running in parallel.

Any new requests get queued and processed as a go routine becomes
free.  The queue is processed in FIFO order.  If a request is
received for an instance that is already in the queue or already
being processed then it will be silently ignored.

*/

package discovery

import (
	"errors"
	"sync"

	"github.com/outbrain/golib/log"
	"github.com/outbrain/orchestrator/go/inst"
)

// Queue is a container for processing the orchestrator discovery requests.
type Queue struct {
	// The number of active go routines processing discovery requests.
	concurrency uint
	// Channel used by the active go routines to say they have finished processing.
	done chan inst.InstanceKey
	// Input channel we are reading the discovery requests from.
	inputChan <-chan inst.InstanceKey
	// This has 2 uses: to indicate if there is a request in the queue
	// (so not being processed) or to indicate that the request is actively
	// being dealth with.  That state is not explicitly stored as it is not
	// really needed.
	knownKeys map[inst.InstanceKey]bool
	// lock while making changes
	lock sync.Mutex
	// The maximum number of go routines allowed to handle the queue at once.
	// This is a configuration parameter provided when creating the Queue.
	maxConcurrency uint
	// process to run on each received key
	processor func(i inst.InstanceKey)
	// This holds the discover requests (the instance name) which need to be
	// processed. but which are not currently being processed. All requests
	// are initially added to the end of this queue, and then the top element
	// will be popped off if the number of active go routines (defined by
	// concurrency) is less than the maximum specified value at which point
	// it will be processed by a new go routine.
	queue []inst.InstanceKey
}

var emptyKey = inst.InstanceKey{}

// NewQueue creates a new Queue entry.
func NewQueue(maxConcurrency uint, inputChan chan inst.InstanceKey, processor func(i inst.InstanceKey)) *Queue {
	log.Infof("Queue.NewQueue()")
	q := new(Queue)

	q.concurrency = 0                    // explicitly
	q.done = make(chan inst.InstanceKey) // Do I need this to be larger?
	q.inputChan = inputChan
	q.knownKeys = make(map[inst.InstanceKey]bool)
	q.maxConcurrency = maxConcurrency
	q.processor = processor
	q.queue = make([]inst.InstanceKey, 0)

	return q
}

// add the key to the slice if it does not exist in known keys
// - goroutine safe as only called inside the mutex
func (q *Queue) push(key inst.InstanceKey) {
	if key == emptyKey {
		log.Fatal("Queue.push(%v) is empty", key)
	}

	if _, found := q.knownKeys[key]; !found {
		// add to the items that are being processed
		q.knownKeys[key] = true
		q.queue = append(q.queue, key)
	} else {
		// If key already there we just ignore it as the request is in the queue.
		// the known key also records stuff in the queue, so pending + active jobs.
	}
}

// remove the entry and remove it from known keys
func (q *Queue) pop() (inst.InstanceKey, error) {
	if len(q.queue) == 0 {
		return inst.InstanceKey{}, errors.New("q.pop() on empty queue")
	}
	key := q.queue[0]
	q.queue = q.queue[1:]
	delete(q.knownKeys, key)
	return key, nil
}

// dispatch a job from the queue (assumes we are in a locked state)
func (q *Queue) dispatch() {
	key, err := q.pop() // should never give an error but let's check anyway
	if err != nil {
		log.Fatal("Queue.dispatch() q.pop() returns: %+v", err)
		return
	}
	if key == emptyKey {
		log.Fatal("Queue.dispatch() key is empty")
	}

	q.concurrency++
	q.knownKeys[key] = true

	// dispatch a discoverInstance() but tell us when we're done (to limit concurrency)
	go func() { // discover asynchronously
		q.processor(key)
		q.done <- key
	}()
}

// acknowledge a job has finished
// - we deal with the locking inside
func (q *Queue) acknowledgeJob(key inst.InstanceKey) {
	q.lock.Lock()
	delete(q.knownKeys, key)
	q.concurrency--
	q.lock.Unlock()
}

// drain queue by dispatching any jobs we have still
func (q *Queue) maybeDispatch() {
	q.lock.Lock()
	if q.concurrency < q.maxConcurrency && len(q.queue) > 0 {
		q.dispatch()
	}
	q.lock.Unlock()
}

// add an entry to the queue and dispatch something if concurrency is low enough
// - we deal with locking inside
func (q *Queue) queueAndMaybeDispatch(key inst.InstanceKey) {
	if key == emptyKey {
		log.Fatal("Queue.queueAndMaybeDispatch(%v) is empty", key)
	}
	q.lock.Lock()
	q.push(key)
	if q.concurrency < q.maxConcurrency && len(q.queue) > 0 {
		q.dispatch()
	}
	q.lock.Unlock()
}

// cleanup is called when the input channel closes.
// we can not sit in the loop so we have to wait for running go-routines to finish
// but also to dispatch anything left in the queue until finally everything is done.
func (q *Queue) cleanup() {
	log.Infof("Queue.cleanup()")
	for q.concurrency > 0 && len(q.queue) > 0 {
		q.maybeDispatch()
		if key, closed := <-q.done; closed {
			return
		} else {
			q.acknowledgeJob(key)
		}
	}
}

// Ends when all elements in the queue have been handled.
// we read from inputChan and call processor up to maxConcurrency times in parallel
func (q *Queue) HandleRequests() {
	if q == nil {
		log.Infof("Queue.HandleRequests() q == nil ??. Should not happen")

		// no queue, nothing to do
		return
	}
	log.Infof("Queue.NewQueue() processing requests")
	for {
		select {
		case key, ok := <-q.inputChan:
			if ok {
				if key != emptyKey {
					q.queueAndMaybeDispatch(key)
				} else {
					log.Warningf("Queue.HandleRequests() q.inputChan received empty key %+v, ignoring (fix the upstream code to prevent this)", key)
				}
			} else {
				q.cleanup()
				log.Infof("Queue.HandleRequests() q.inputChan is closed. returning")
				return
			}
		case key, ok := <-q.done:
			if ok {
				q.acknowledgeJob(key)
			} else {
				log.Infof("Queue.HandleRequests() q.done is closed. returning (shouldn't get here)")
				return // we shouldn't get here as the return above should get triggered first
			}
		}
	}
}
<|endoftext|>


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

package search

import (
	"context"
	"os"
	"path/filepath"
	"sync"

	"github.com/blevesearch/bleve"
	"github.com/blevesearch/bleve/index/store"
	"github.com/blevesearch/bleve/mapping"
	"github.com/blevesearch/bleve/registry"
	"github.com/keybase/client/go/kbfs/data"
	"github.com/keybase/client/go/kbfs/idutil"
	"github.com/keybase/client/go/kbfs/kbfsmd"
	"github.com/keybase/client/go/kbfs/libfs"
	"github.com/keybase/client/go/kbfs/libkbfs"
	"github.com/keybase/client/go/kbfs/tlf"
	kbname "github.com/keybase/client/go/kbun"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/pkg/errors"
)

const (
	textFileType      = "kbfsTextFile"
	kvstoreName       = "kbfs"
	bleveIndexType    = "upside_down"
	fsIndexStorageDir = "kbfs_index"
)

const (
	// CtxOpID is the display name for the unique operation index ID tag.
	ctxOpID = "IID"
)

// CtxTagKey is the type used for unique context tags
type ctxTagKey int

const (
	// CtxIDKey is the type of the tag for unique operation IDs.
	ctxIDKey ctxTagKey = iota
)

type tlfMessage struct {
	tlfID tlf.ID
	rev   kbfsmd.Revision
	mode  keybase1.FolderSyncMode
}

// Indexer can index and search KBFS TLFs.
type Indexer struct {
	config       libkbfs.Config
	log          logger.Logger
	cancelLoop   context.CancelFunc
	remoteStatus libfs.RemoteStatus

	userChangedCh chan struct{}
	tlfCh         chan tlfMessage
	shutdownCh    chan struct{}

	lock        sync.RWMutex
	index       bleve.Index
	indexConfig libkbfs.Config
}

// NewIndexer creates a new instance of an Indexer.
func NewIndexer(config libkbfs.Config) (*Indexer, error) {
	log := config.MakeLogger("search")
	i := &Indexer{
		config:        config,
		log:           log,
		userChangedCh: make(chan struct{}, 1),
		tlfCh:         make(chan tlfMessage, 1000),
		shutdownCh:    make(chan struct{}),
	}

	ctx, cancel := context.WithCancel(i.makeContext(context.Background()))
	i.cancelLoop = cancel
	go i.loop(ctx)
	err := config.Notifier().RegisterForSyncedTlfs(i)
	if err != nil {
		return nil, err
	}
	return i, nil
}

func (i *Indexer) makeContext(ctx context.Context) context.Context {
	return libkbfs.CtxWithRandomIDReplayable(ctx, ctxIDKey, ctxOpID, i.log)
}

func (i *Indexer) closeIndexLocked(ctx context.Context) error {
	if i.index == nil {
		return nil
	}

	err := i.index.Close()
	if err != nil {
		return err
	}
	i.index = nil

	return i.indexConfig.Shutdown(ctx)
}

func (i *Indexer) loadIndex(ctx context.Context) (err error) {
	i.lock.Lock()
	defer i.lock.Unlock()

	err = i.closeIndexLocked(ctx)
	if err != nil {
		return err
	}

	session, err := idutil.GetCurrentSessionIfPossible(
		ctx, i.config.KBPKI(), true)
	if err != nil {
		return err
	}
	if session.Name == "" {
		i.log.CDebugf(ctx, "Not indexing while logged out")
		return nil
	}

	// Create a new Config object for the index data, with a storage
	// root that's unique to this user.
	kbCtx := i.config.KbContext()
	params, err := Params(kbCtx, i.config.StorageRoot(), session.UID)
	if err != nil {
		return err
	}
	ctx, indexConfig, err := Init(
		ctx, kbCtx, params, libkbfs.NewKeybaseServicePassthrough(i.config),
		i.log, i.config.VLogLevel())
	if err != nil {
		return err
	}
	defer func() {
		if err != nil {
			configErr := indexConfig.Shutdown(ctx)
			if configErr != nil {
				i.log.CDebugf(ctx, "Couldn't shutdown config: %+v", configErr)
			}
		}
	}()

	// Store the index in a KBFS private folder for the current user,
	// with all the blocks and MD stored in the storage root created
	// above.  Everything will be encrypted as if it were in the
	// user's own private KBFS folder.
	privateHandle, err := libkbfs.GetHandleFromFolderNameAndType(
		ctx, indexConfig.KBPKI(), indexConfig.MDOps(), indexConfig,
		string(session.Name), tlf.Private)
	if err != nil {
		return err
	}
	fs, err := libfs.NewFS(
		ctx, indexConfig, privateHandle, data.MasterBranch, "", "", 0)
	if err != nil {
		return err
	}
	err = fs.MkdirAll(fsIndexStorageDir, 0400)
	if err != nil {
		return err
	}
	fs, err = fs.ChrootAsLibFS(fsIndexStorageDir)
	if err != nil {
		return err
	}

	// The index itself will use LevelDB storage that writes to the
	// KBFS filesystem object made above.  Register this storage type
	// with Bleve.
	kvstoreConstructor := func(
		mo store.MergeOperator, _ map[string]interface{}) (
		s store.KVStore, err error) {
		return newBleveLevelDBStore(fs, false, mo)
	}
	registry.RegisterKVStore(kvstoreName, kvstoreConstructor)

	// Create the actual index using this storage type.  Bleve has
	// different calls for new vs. existing indicies, so we first need
	// to check if it exists.  The Bleve LevelDB storage takes a lock,
	// so we don't really need to worry about concurrent KBFS
	// processes here.
	var index bleve.Index
	p := filepath.Join(i.config.StorageRoot(), indexStorageDir, bleveIndexDir)
	_, err = os.Stat(p)
	switch {
	case os.IsNotExist(errors.Cause(err)):
		i.log.CDebugf(ctx, "Creating new index for user %s/%s",
			session.Name, session.UID)

		// TODO: add an html/md document parser to get rid of tags.
		indexMapping := bleve.NewIndexMapping()
		m := mapping.NewDocumentMapping()
		indexMapping.AddDocumentMapping(textFileType, m)
		index, err = bleve.NewUsing(
			p, indexMapping, bleveIndexType, kvstoreName, nil)
		if err != nil {
			return err
		}
	case err == nil:
		i.log.CDebugf(ctx, "Using existing index for user %s/%s",
			session.Name, session.UID)

		index, err = bleve.OpenUsing(p, nil)
		if err != nil {
			return err
		}
	default:
		return err
	}

	i.index = index
	i.indexConfig = indexConfig
	return nil
}

// UserChanged implements the libfs.RemoteStatusUpdater for Indexer.
func (i *Indexer) UserChanged(
	ctx context.Context, oldName, newName kbname.NormalizedUsername) {
	select {
	case i.userChangedCh <- struct{}{}:
	default:
		i.log.CDebugf(ctx, "Dropping user changed notification")
	}
}

var _ libfs.RemoteStatusUpdater = (*Indexer)(nil)

// FullSyncStarted implements the libkbfs.SyncedTlfObserver interface
// for Indexer.
func (i *Indexer) FullSyncStarted(
	ctx context.Context, tlfID tlf.ID, rev kbfsmd.Revision,
	waitCh <-chan struct{}) {
	i.log.CDebugf(ctx, "Sync started for %s/%d", tlfID, rev)
	go func() {
		select {
		case <-waitCh:
		case <-i.shutdownCh:
			return
		}

		m := tlfMessage{tlfID, rev, keybase1.FolderSyncMode_ENABLED}
		select {
		case i.tlfCh <- m:
		default:
			i.log.CDebugf(
				context.Background(), "Couldn't send TLF message for %s/%d")
		}
	}()
}

// SyncModeChanged implements the libkbfs.SyncedTlfObserver interface
// for Indexer.
func (i *Indexer) SyncModeChanged(
	ctx context.Context, tlfID tlf.ID, newMode keybase1.FolderSyncMode) {
	i.log.CDebugf(ctx, "Sync mode changed for %s to %s", tlfID, newMode)
	m := tlfMessage{tlfID, kbfsmd.RevisionUninitialized, newMode}
	select {
	case i.tlfCh <- m:
	default:
		i.log.CDebugf(
			context.Background(), "Couldn't send TLF message for %s/%d")
	}
}

var _ libkbfs.SyncedTlfObserver = (*Indexer)(nil)

func (i *Indexer) loop(ctx context.Context) {
	i.log.CDebugf(ctx, "Starting indexing loop")
	defer i.log.CDebugf(ctx, "Ending index loop")

	i.remoteStatus.Init(ctx, i.log, i.config, i)

outerLoop:
	for {
		err := i.loadIndex(ctx)
		if err != nil {
			i.log.CDebugf(ctx, "Couldn't load index: %+v", err)
		}

		state := keybase1.MobileAppState_FOREGROUND
		kbCtx := i.config.KbContext()
		for {
			select {
			case <-i.userChangedCh:
				// Re-load the index on each login/logout event.
				i.log.CDebugf(ctx, "User changed")
				continue outerLoop
			case state = <-kbCtx.NextAppStateUpdate(&state):
				// TODO(HOTPOT-1494): once we are doing actual
				// indexing in a separate goroutine, pause/unpause it
				// via a channel send from here.
				for state != keybase1.MobileAppState_FOREGROUND {
					i.log.CDebugf(ctx,
						"Pausing indexing while not foregrounded: state=%s",
						state)
					state = <-kbCtx.NextAppStateUpdate(&state)
				}
				i.log.CDebugf(ctx, "Resuming indexing while foregrounded")
				continue
			case m := <-i.tlfCh:
				i.log.CDebugf(ctx, "Received TLF message for %s", m.tlfID)
				// TODO(HOTPOT-1494, HOTPOT-1495): initiate processing pass.
			case <-ctx.Done():
				return
			}
		}
	}
}

// Shutdown shuts down this indexer.
func (i *Indexer) Shutdown(ctx context.Context) error {
	i.cancelLoop()
	close(i.shutdownCh)

	i.lock.Lock()
	defer i.lock.Unlock()

	return i.closeIndexLocked(ctx)
}
<commit_msg>search: make an FS locked at a revision for indexing<commit_after>// Copyright 2019 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package search

import (
	"context"
	"os"
	"path/filepath"
	"sync"

	"github.com/blevesearch/bleve"
	"github.com/blevesearch/bleve/index/store"
	"github.com/blevesearch/bleve/mapping"
	"github.com/blevesearch/bleve/registry"
	"github.com/keybase/client/go/kbfs/data"
	"github.com/keybase/client/go/kbfs/idutil"
	"github.com/keybase/client/go/kbfs/kbfsmd"
	"github.com/keybase/client/go/kbfs/libfs"
	"github.com/keybase/client/go/kbfs/libkbfs"
	"github.com/keybase/client/go/kbfs/tlf"
	kbname "github.com/keybase/client/go/kbun"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/pkg/errors"
)

const (
	textFileType      = "kbfsTextFile"
	kvstoreName       = "kbfs"
	bleveIndexType    = "upside_down"
	fsIndexStorageDir = "kbfs_index"
)

const (
	// CtxOpID is the display name for the unique operation index ID tag.
	ctxOpID = "IID"
)

// CtxTagKey is the type used for unique context tags
type ctxTagKey int

const (
	// CtxIDKey is the type of the tag for unique operation IDs.
	ctxIDKey ctxTagKey = iota
)

type tlfMessage struct {
	tlfID tlf.ID
	rev   kbfsmd.Revision
	mode  keybase1.FolderSyncMode
}

// Indexer can index and search KBFS TLFs.
type Indexer struct {
	config       libkbfs.Config
	log          logger.Logger
	cancelLoop   context.CancelFunc
	remoteStatus libfs.RemoteStatus

	userChangedCh chan struct{}
	tlfCh         chan tlfMessage
	shutdownCh    chan struct{}

	lock        sync.RWMutex
	index       bleve.Index
	indexConfig libkbfs.Config
}

// NewIndexer creates a new instance of an Indexer.
func NewIndexer(config libkbfs.Config) (*Indexer, error) {
	log := config.MakeLogger("search")
	i := &Indexer{
		config:        config,
		log:           log,
		userChangedCh: make(chan struct{}, 1),
		tlfCh:         make(chan tlfMessage, 1000),
		shutdownCh:    make(chan struct{}),
	}

	ctx, cancel := context.WithCancel(i.makeContext(context.Background()))
	i.cancelLoop = cancel
	go i.loop(ctx)
	err := config.Notifier().RegisterForSyncedTlfs(i)
	if err != nil {
		return nil, err
	}
	return i, nil
}

func (i *Indexer) makeContext(ctx context.Context) context.Context {
	return libkbfs.CtxWithRandomIDReplayable(ctx, ctxIDKey, ctxOpID, i.log)
}

func (i *Indexer) closeIndexLocked(ctx context.Context) error {
	if i.index == nil {
		return nil
	}

	err := i.index.Close()
	if err != nil {
		return err
	}
	i.index = nil

	return i.indexConfig.Shutdown(ctx)
}

func (i *Indexer) loadIndex(ctx context.Context) (err error) {
	i.lock.Lock()
	defer i.lock.Unlock()

	err = i.closeIndexLocked(ctx)
	if err != nil {
		return err
	}

	session, err := idutil.GetCurrentSessionIfPossible(
		ctx, i.config.KBPKI(), true)
	if err != nil {
		return err
	}
	if session.Name == "" {
		i.log.CDebugf(ctx, "Not indexing while logged out")
		return nil
	}

	// Create a new Config object for the index data, with a storage
	// root that's unique to this user.
	kbCtx := i.config.KbContext()
	params, err := Params(kbCtx, i.config.StorageRoot(), session.UID)
	if err != nil {
		return err
	}
	ctx, indexConfig, err := Init(
		ctx, kbCtx, params, libkbfs.NewKeybaseServicePassthrough(i.config),
		i.log, i.config.VLogLevel())
	if err != nil {
		return err
	}
	defer func() {
		if err != nil {
			configErr := indexConfig.Shutdown(ctx)
			if configErr != nil {
				i.log.CDebugf(ctx, "Couldn't shutdown config: %+v", configErr)
			}
		}
	}()

	// Store the index in a KBFS private folder for the current user,
	// with all the blocks and MD stored in the storage root created
	// above.  Everything will be encrypted as if it were in the
	// user's own private KBFS folder.
	privateHandle, err := libkbfs.GetHandleFromFolderNameAndType(
		ctx, indexConfig.KBPKI(), indexConfig.MDOps(), indexConfig,
		string(session.Name), tlf.Private)
	if err != nil {
		return err
	}
	fs, err := libfs.NewFS(
		ctx, indexConfig, privateHandle, data.MasterBranch, "", "", 0)
	if err != nil {
		return err
	}
	err = fs.MkdirAll(fsIndexStorageDir, 0400)
	if err != nil {
		return err
	}
	fs, err = fs.ChrootAsLibFS(fsIndexStorageDir)
	if err != nil {
		return err
	}

	// The index itself will use LevelDB storage that writes to the
	// KBFS filesystem object made above.  Register this storage type
	// with Bleve.
	kvstoreConstructor := func(
		mo store.MergeOperator, _ map[string]interface{}) (
		s store.KVStore, err error) {
		return newBleveLevelDBStore(fs, false, mo)
	}
	registry.RegisterKVStore(kvstoreName, kvstoreConstructor)

	// Create the actual index using this storage type.  Bleve has
	// different calls for new vs. existing indicies, so we first need
	// to check if it exists.  The Bleve LevelDB storage takes a lock,
	// so we don't really need to worry about concurrent KBFS
	// processes here.
	var index bleve.Index
	p := filepath.Join(i.config.StorageRoot(), indexStorageDir, bleveIndexDir)
	_, err = os.Stat(p)
	switch {
	case os.IsNotExist(errors.Cause(err)):
		i.log.CDebugf(ctx, "Creating new index for user %s/%s",
			session.Name, session.UID)

		// Register a mapping for text files, so when we index text
		// files we can mark them as such. TODO(HOTPOT-1488): add an
		// html/md document parser to get rid of tags.
		indexMapping := bleve.NewIndexMapping()
		m := mapping.NewDocumentMapping()
		indexMapping.AddDocumentMapping(textFileType, m)
		index, err = bleve.NewUsing(
			p, indexMapping, bleveIndexType, kvstoreName, nil)
		if err != nil {
			return err
		}
	case err == nil:
		i.log.CDebugf(ctx, "Using existing index for user %s/%s",
			session.Name, session.UID)

		index, err = bleve.OpenUsing(p, nil)
		if err != nil {
			return err
		}
	default:
		return err
	}

	i.index = index
	i.indexConfig = indexConfig
	return nil
}

// UserChanged implements the libfs.RemoteStatusUpdater for Indexer.
func (i *Indexer) UserChanged(
	ctx context.Context, oldName, newName kbname.NormalizedUsername) {
	select {
	case i.userChangedCh <- struct{}{}:
	default:
		i.log.CDebugf(ctx, "Dropping user changed notification")
	}
}

var _ libfs.RemoteStatusUpdater = (*Indexer)(nil)

// FullSyncStarted implements the libkbfs.SyncedTlfObserver interface
// for Indexer.
func (i *Indexer) FullSyncStarted(
	ctx context.Context, tlfID tlf.ID, rev kbfsmd.Revision,
	waitCh <-chan struct{}) {
	i.log.CDebugf(ctx, "Sync started for %s/%d", tlfID, rev)
	go func() {
		select {
		case <-waitCh:
		case <-i.shutdownCh:
			return
		}

		m := tlfMessage{tlfID, rev, keybase1.FolderSyncMode_ENABLED}
		select {
		case i.tlfCh <- m:
		default:
			i.log.CDebugf(
				context.Background(), "Couldn't send TLF message for %s/%d")
		}
	}()
}

// SyncModeChanged implements the libkbfs.SyncedTlfObserver interface
// for Indexer.
func (i *Indexer) SyncModeChanged(
	ctx context.Context, tlfID tlf.ID, newMode keybase1.FolderSyncMode) {
	i.log.CDebugf(ctx, "Sync mode changed for %s to %s", tlfID, newMode)
	m := tlfMessage{tlfID, kbfsmd.RevisionUninitialized, newMode}
	select {
	case i.tlfCh <- m:
	default:
		i.log.CDebugf(
			context.Background(), "Couldn't send TLF message for %s/%d")
	}
}

var _ libkbfs.SyncedTlfObserver = (*Indexer)(nil)

func (i *Indexer) fsForRev(
	ctx context.Context, tlfID tlf.ID, rev kbfsmd.Revision) (*libfs.FS, error) {
	if rev == kbfsmd.RevisionUninitialized {
		return nil, errors.New("No revision provided")
	}
	branch := data.MakeRevBranchName(rev)

	md, err := i.config.MDOps().GetForTLF(ctx, tlfID, nil)
	if err != nil {
		return nil, err
	}

	h := md.GetTlfHandle()
	return libfs.NewReadonlyFS(
		ctx, i.config, h, branch, "", "", keybase1.MDPriorityNormal)
}

func (i *Indexer) loop(ctx context.Context) {
	i.log.CDebugf(ctx, "Starting indexing loop")
	defer i.log.CDebugf(ctx, "Ending index loop")

	i.remoteStatus.Init(ctx, i.log, i.config, i)

outerLoop:
	for {
		err := i.loadIndex(ctx)
		if err != nil {
			i.log.CDebugf(ctx, "Couldn't load index: %+v", err)
		}

		state := keybase1.MobileAppState_FOREGROUND
		kbCtx := i.config.KbContext()
		for {
			select {
			case <-i.userChangedCh:
				// Re-load the index on each login/logout event.
				i.log.CDebugf(ctx, "User changed")
				continue outerLoop
			case state = <-kbCtx.NextAppStateUpdate(&state):
				// TODO(HOTPOT-1494): once we are doing actual
				// indexing in a separate goroutine, pause/unpause it
				// via a channel send from here.
				for state != keybase1.MobileAppState_FOREGROUND {
					i.log.CDebugf(ctx,
						"Pausing indexing while not foregrounded: state=%s",
						state)
					state = <-kbCtx.NextAppStateUpdate(&state)
				}
				i.log.CDebugf(ctx, "Resuming indexing while foregrounded")
				continue
			case m := <-i.tlfCh:
				ctx := i.makeContext(ctx)
				i.log.CDebugf(ctx, "Received TLF message for %s", m.tlfID)

				// Figure out which revision to lock to, for this
				// indexing scan.
				var rev kbfsmd.Revision
				if m.rev == kbfsmd.RevisionUninitialized {
					// TODO(HOTPOT-1504) -- remove indexing if the
					// mode is no longer synced.
					continue
				}

				_, err := i.fsForRev(ctx, m.tlfID, rev)
				if err != nil {
					i.log.CDebugf(ctx, "Error making FS: %+v", err)
					continue
				}

				// TODO(HOTPOT-1494, HOTPOT-1495): initiate processing
				// pass, using the FS made above.
			case <-ctx.Done():
				return
			}
		}
	}
}

// Shutdown shuts down this indexer.
func (i *Indexer) Shutdown(ctx context.Context) error {
	i.cancelLoop()
	close(i.shutdownCh)

	i.lock.Lock()
	defer i.lock.Unlock()

	return i.closeIndexLocked(ctx)
}
<|endoftext|>


<commit_before>package virt

import (
	"errors"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"labix.org/v2/mgo/bson"
)

func (vm *VM) Start() error {
	if out, err := exec.Command("/usr/bin/lxc-start", "--name", vm.String(), "--daemon").CombinedOutput(); err != nil {
		return commandError("lxc-start failed.", err, out)
	}
	return vm.WaitForState("RUNNING", time.Second)
}

func (vm *VM) Stop() error {
	if out, err := exec.Command("/usr/bin/lxc-stop", "--name", vm.String()).CombinedOutput(); err != nil {
		return commandError("lxc-stop failed.", err, out)
	}
	return vm.WaitForState("STOPPED", time.Second)
}

func (vm *VM) Shutdown() error {
	if out, err := exec.Command("/usr/bin/lxc-shutdown", "--name", vm.String()).CombinedOutput(); err != nil {
		if vm.GetState() != "STOPPED" {
			return commandError("lxc-shutdown failed.", err, out)
		}
	}
	vm.WaitForState("STOPPED", 5*time.Second) // may time out, then vm is force stopped
	return vm.Stop()
}

func (vm *VM) AttachCommand(uid int, tty string, command ...string) *exec.Cmd {
	args := []string{"--name", vm.String()}
	if tty != "" {
		args = append(args, "--tty", tty)
	}
	args = append(args, "--", "/usr/bin/sudo", "-i", "-u", "#"+strconv.Itoa(uid), "--")
	args = append(args, command...)
	cmd := exec.Command("/usr/bin/lxc-attach", args...)
	cmd.Env = []string{"TERM=xterm-256color"}
	return cmd
}

func GetVMState(vmId bson.ObjectId) string {
	out, err := exec.Command("/usr/bin/lxc-info", "--name", VMName(vmId), "--state").CombinedOutput()
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(out)[6:])
}

func (vm *VM) GetState() string {
	return GetVMState(vm.Id)
}

func (vm *VM) WaitForState(state string, timeout time.Duration) error {
	tryUntil := time.Now().Add(timeout)
	for vm.GetState() != state {
		if time.Now().After(tryUntil) {
			return errors.New("Timeout while waiting for VM state.")
		}
		time.Sleep(time.Second / 10)
	}
	return nil
}

func (vm *VM) SendMessageToVMUsers(message string) error {
	cmd := exec.Command("/usr/bin/lxc-attach", "--name", vm.String(), "--", "/usr/bin/wall", "--nobanner")
	cmd.Stdin = strings.NewReader(message)
	if out, err := cmd.CombinedOutput(); err != nil {
		return commandError("wall failed.", err, out)
	}
	return nil
}

func (vm *VM) WaitForNetwork(timeout time.Duration) error {
	// precaution for bad input
	if timeout == 0 {
		timeout = time.Second * 5
	}

	isNetworkUp := func() bool {
		// neglect error because it's not important and we also going to timeout
		out, _ := exec.Command("/usr/bin/lxc-attach", "--name", vm.String(),
			"--", "/bin/cat", "/sys/class/net/eth0/operstate").CombinedOutput()

		if strings.TrimSpace(string(out)) == "up" {
			return true

		return false
	}

	tryUntil := time.Now().Add(timeout)
	for {
		if up := isNetworkUp(); up {
			return nil
		}

		if time.Now().After(tryUntil) {
			return errors.New("Timeout while waiting for VM Network state.")
		}

		time.Sleep(time.Millisecond * 100)
	}
}
<commit_msg>oskite: oops, saved it accidently :)<commit_after>package virt

import (
	"errors"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"labix.org/v2/mgo/bson"
)

func (vm *VM) Start() error {
	if out, err := exec.Command("/usr/bin/lxc-start", "--name", vm.String(), "--daemon").CombinedOutput(); err != nil {
		return commandError("lxc-start failed.", err, out)
	}
	return vm.WaitForState("RUNNING", time.Second)
}

func (vm *VM) Stop() error {
	if out, err := exec.Command("/usr/bin/lxc-stop", "--name", vm.String()).CombinedOutput(); err != nil {
		return commandError("lxc-stop failed.", err, out)
	}
	return vm.WaitForState("STOPPED", time.Second)
}

func (vm *VM) Shutdown() error {
	if out, err := exec.Command("/usr/bin/lxc-shutdown", "--name", vm.String()).CombinedOutput(); err != nil {
		if vm.GetState() != "STOPPED" {
			return commandError("lxc-shutdown failed.", err, out)
		}
	}
	vm.WaitForState("STOPPED", 5*time.Second) // may time out, then vm is force stopped
	return vm.Stop()
}

func (vm *VM) AttachCommand(uid int, tty string, command ...string) *exec.Cmd {
	args := []string{"--name", vm.String()}
	if tty != "" {
		args = append(args, "--tty", tty)
	}
	args = append(args, "--", "/usr/bin/sudo", "-i", "-u", "#"+strconv.Itoa(uid), "--")
	args = append(args, command...)
	cmd := exec.Command("/usr/bin/lxc-attach", args...)
	cmd.Env = []string{"TERM=xterm-256color"}
	return cmd
}

func GetVMState(vmId bson.ObjectId) string {
	out, err := exec.Command("/usr/bin/lxc-info", "--name", VMName(vmId), "--state").CombinedOutput()
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(out)[6:])
}

func (vm *VM) GetState() string {
	return GetVMState(vm.Id)
}

func (vm *VM) WaitForState(state string, timeout time.Duration) error {
	tryUntil := time.Now().Add(timeout)
	for vm.GetState() != state {
		if time.Now().After(tryUntil) {
			return errors.New("Timeout while waiting for VM state.")
		}
		time.Sleep(time.Second / 10)
	}
	return nil
}

func (vm *VM) SendMessageToVMUsers(message string) error {
	cmd := exec.Command("/usr/bin/lxc-attach", "--name", vm.String(), "--", "/usr/bin/wall", "--nobanner")
	cmd.Stdin = strings.NewReader(message)
	if out, err := cmd.CombinedOutput(); err != nil {
		return commandError("wall failed.", err, out)
	}
	return nil
}

func (vm *VM) WaitForNetwork(timeout time.Duration) error {
	// precaution for bad input
	if timeout == 0 {
		timeout = time.Second * 5
	}

	isNetworkUp := func() bool {
		// neglect error because it's not important and we also going to timeout
		out, _ := exec.Command("/usr/bin/lxc-attach", "--name", vm.String(),
			"--", "/bin/cat", "/sys/class/net/eth0/operstate").CombinedOutput()

		if strings.TrimSpace(string(out)) == "up" {
			return true
		}

		return false
	}

	tryUntil := time.Now().Add(timeout)
	for {
		if up := isNetworkUp(); up {
			return nil
		}

		if time.Now().After(tryUntil) {
			return errors.New("Timeout while waiting for VM Network state.")
		}

		time.Sleep(time.Millisecond * 100)
	}
}
<|endoftext|>


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

package stats

import (
	"expvar"
	"reflect"
	"sort"
	"testing"
	"time"
)

func TestCounters(t *testing.T) {
	clear()
	c := NewCounters("counter1")
	c.Add("c1", 1)
	c.Add("c2", 1)
	c.Add("c2", 1)
	want1 := `{"c1": 1, "c2": 2}`
	want2 := `{"c2": 2, "c1": 1}`
	if s := c.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
	counts := c.Counts()
	if counts["c1"] != 1 {
		t.Errorf("want 1, got %d", counts["c1"])
	}
	if counts["c2"] != 2 {
		t.Errorf("want 2, got %d", counts["c2"])
	}
	f := CountersFunc(func() map[string]int64 {
		return map[string]int64{
			"c1": 1,
			"c2": 2,
		}
	})
	if s := f.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
}

func TestCountersTags(t *testing.T) {
	clear()
	c := NewCounters("counterTag1")
	want := map[string]int64{}
	got := c.Counts()
	if !reflect.DeepEqual(got, want) {
		t.Errorf("want %v, got %v", want, got)
	}

	c = NewCounters("counterTag2", "tag1", "tag2")
	want = map[string]int64{"tag1": 0, "tag2": 0}
	got = c.Counts()
	if !reflect.DeepEqual(got, want) {
		t.Errorf("want %v, got %v", want, got)
	}
}

func TestMultiCounters(t *testing.T) {
	clear()
	c := NewMultiCounters("mapCounter1", []string{"aaa", "bbb"})
	c.Add([]string{"c1a", "c1b"}, 1)
	c.Add([]string{"c2a", "c2b"}, 1)
	c.Add([]string{"c2a", "c2b"}, 1)
	want1 := `{"c1a.c1b": 1, "c2a.c2b": 2}`
	want2 := `{"c2a.c2b": 2, "c1a.c1b": 1}`
	if s := c.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
	counts := c.Counts()
	if counts["c1a.c1b"] != 1 {
		t.Errorf("want 1, got %d", counts["c1a.c1b"])
	}
	if counts["c2a.c2b"] != 2 {
		t.Errorf("want 2, got %d", counts["c2a.c2b"])
	}
	f := NewMultiCountersFunc("", []string{"aaa", "bbb"}, func() map[string]int64 {
		return map[string]int64{
			"c1a.c1b": 1,
			"c2a.c2b": 2,
		}
	})
	if s := f.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
}

func TestCountersHook(t *testing.T) {
	var gotname string
	var gotv *Counters
	clear()
	Register(func(name string, v expvar.Var) {
		gotname = name
		gotv = v.(*Counters)
	})

	v := NewCounters("counter2")
	if gotname != "counter2" {
		t.Errorf("want counter2, got %s", gotname)
	}
	if gotv != v {
		t.Errorf("want %#v, got %#v", v, gotv)
	}
}

var benchCounter = NewCounters("bench")

func BenchmarkCounters(b *testing.B) {
	clear()
	benchCounter.Add("c1", 1)
	b.ResetTimer()

	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			benchCounter.Add("c1", 1)
		}
	})
}

func BenchmarkCountersTailLatency(b *testing.B) {
	// For this one, ignore the time reported by 'go test'.
	// The 99th Percentile log line is all that matters.
	clear()
	benchCounter.Add("c1", 1)
	c := make(chan time.Duration, 100)
	done := make(chan struct{})
	go func() {
		all := make([]int, b.N)
		i := 0
		for dur := range c {
			all[i] = int(dur)
			i++
		}
		sort.Ints(all)
		p99 := time.Duration(all[b.N*99/100])
		b.Logf("99th Percentile (for N=%v): %v", b.N, p99)
		close(done)
	}()

	b.ResetTimer()
	b.SetParallelism(1000)
	b.RunParallel(func(pb *testing.PB) {
		var start time.Time

		for pb.Next() {
			start = time.Now()
			benchCounter.Add("c1", 1)
			c <- time.Since(start)
		}
	})
	b.StopTimer()

	close(c)
	<-done
}
<commit_msg>Improve counters 99th latency benchmark.<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package stats

import (
	"expvar"
	"math/rand"
	"reflect"
	"sort"
	"testing"
	"time"
)

func TestCounters(t *testing.T) {
	clear()
	c := NewCounters("counter1")
	c.Add("c1", 1)
	c.Add("c2", 1)
	c.Add("c2", 1)
	want1 := `{"c1": 1, "c2": 2}`
	want2 := `{"c2": 2, "c1": 1}`
	if s := c.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
	counts := c.Counts()
	if counts["c1"] != 1 {
		t.Errorf("want 1, got %d", counts["c1"])
	}
	if counts["c2"] != 2 {
		t.Errorf("want 2, got %d", counts["c2"])
	}
	f := CountersFunc(func() map[string]int64 {
		return map[string]int64{
			"c1": 1,
			"c2": 2,
		}
	})
	if s := f.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
}

func TestCountersTags(t *testing.T) {
	clear()
	c := NewCounters("counterTag1")
	want := map[string]int64{}
	got := c.Counts()
	if !reflect.DeepEqual(got, want) {
		t.Errorf("want %v, got %v", want, got)
	}

	c = NewCounters("counterTag2", "tag1", "tag2")
	want = map[string]int64{"tag1": 0, "tag2": 0}
	got = c.Counts()
	if !reflect.DeepEqual(got, want) {
		t.Errorf("want %v, got %v", want, got)
	}
}

func TestMultiCounters(t *testing.T) {
	clear()
	c := NewMultiCounters("mapCounter1", []string{"aaa", "bbb"})
	c.Add([]string{"c1a", "c1b"}, 1)
	c.Add([]string{"c2a", "c2b"}, 1)
	c.Add([]string{"c2a", "c2b"}, 1)
	want1 := `{"c1a.c1b": 1, "c2a.c2b": 2}`
	want2 := `{"c2a.c2b": 2, "c1a.c1b": 1}`
	if s := c.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
	counts := c.Counts()
	if counts["c1a.c1b"] != 1 {
		t.Errorf("want 1, got %d", counts["c1a.c1b"])
	}
	if counts["c2a.c2b"] != 2 {
		t.Errorf("want 2, got %d", counts["c2a.c2b"])
	}
	f := NewMultiCountersFunc("", []string{"aaa", "bbb"}, func() map[string]int64 {
		return map[string]int64{
			"c1a.c1b": 1,
			"c2a.c2b": 2,
		}
	})
	if s := f.String(); s != want1 && s != want2 {
		t.Errorf("want %s or %s, got %s", want1, want2, s)
	}
}

func TestCountersHook(t *testing.T) {
	var gotname string
	var gotv *Counters
	clear()
	Register(func(name string, v expvar.Var) {
		gotname = name
		gotv = v.(*Counters)
	})

	v := NewCounters("counter2")
	if gotname != "counter2" {
		t.Errorf("want counter2, got %s", gotname)
	}
	if gotv != v {
		t.Errorf("want %#v, got %#v", v, gotv)
	}
}

var benchCounter = NewCounters("bench")

func BenchmarkCounters(b *testing.B) {
	clear()
	benchCounter.Add("c1", 1)
	b.ResetTimer()

	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			benchCounter.Add("c1", 1)
		}
	})
}

func BenchmarkCountersTailLatency(b *testing.B) {
	// For this one, ignore the time reported by 'go test'.
	// The 99th Percentile log line is all that matters.
	// (Cmd: go test -bench=BenchmarkCountersTailLatency -benchtime=30s -cpu=10)
	clear()
	benchCounter.Add("c1", 1)
	c := make(chan time.Duration, 100)
	done := make(chan struct{})
	go func() {
		all := make([]int, b.N)
		i := 0
		for dur := range c {
			all[i] = int(dur)
			i++
		}
		sort.Ints(all)
		p99 := time.Duration(all[b.N*99/100])
		b.Logf("99th Percentile (for N=%v): %v", b.N, p99)
		close(done)
	}()

	b.ResetTimer()
	b.SetParallelism(100) // The actual number of goroutines is 100*GOMAXPROCS
	b.RunParallel(func(pb *testing.PB) {
		r := rand.New(rand.NewSource(time.Now().UnixNano()))

		var start time.Time

		for pb.Next() {
			// sleep between 0~200ms to simulate 10 QPS per goroutine.
			time.Sleep(time.Duration(r.Int63n(200)) * time.Millisecond)
			start = time.Now()
			benchCounter.Add("c1", 1)
			c <- time.Since(start)
		}
	})
	b.StopTimer()

	close(c)
	<-done
}
<|endoftext|>


<commit_before>package goarmorapi

import "github.com/pkg/errors"

type Conformance struct {
	Name   string             `json:"name"`
	Server *ConformanceServer `json:"server"`
	API    *ConformanceClient `json:"api"`
}

type ConformanceServer struct {
	Type         string `json:"type"`
	ID           uint64 `json:"id"`
	Version      uint64 `json:"version"`
	ReleaseStage string `json:"releaseStage"`
	Sharded      bool   `json:"sharded,omitempty"`
	ShardsCount  uint64 `json:"shardsCount,omitempty"`
	StartedAt    string `json:"startedAt,omitempty"`
	Architecture string `json:"architecture"`
}

type ConformanceClient struct {
	Version string `json:"version"`
}

func NewConformance(
	apiVersion string,
	infrastructureVersion uint64,
	releaseStageName string,
	serverTitle, serverName, serverArchitecture string,
	serverID uint64,
	shardsCount int,
	startedAt string) (*Conformance, error) {
	if shardsCount < 0 {
		return nil, errors.New("unexpected shard servers count")
	}

	v := &Conformance{
		Name: serverTitle,
		API:  &ConformanceClient{Version: apiVersion},
		Server: &ConformanceServer{
			Type:         serverName,
			ID:           serverID,
			Version:      infrastructureVersion,
			ReleaseStage: releaseStageName,
			Sharded:      shardsCount > 0,
			ShardsCount:  uint64(shardsCount),
			Architecture: serverArchitecture,
			StartedAt:    startedAt}}

	return v, nil
}
<commit_msg>temporary commit<commit_after>package goarmorapi

import "github.com/pkg/errors"

type Conformance struct {
	Name   string             `json:"name"`
	Server *ConformanceServer `json:"server"`
	API    *ConformanceClient `json:"api"`
}

type ConformanceServer struct {
	Type                    string `json:"type"`
	ID                      uint64 `json:"id"`
	Version                 uint64 `json:"version"`
	ReleaseStage            string `json:"releaseStage"`
	Sharded                 bool   `json:"sharded,omitempty"`
	ShardsCount             uint64 `json:"shardsCount,omitempty"`
	StartedAt               string `json:"startedAt,omitempty"`
	Architecture            string `json:"architecture"`
	HeartbeatMetricsEnabled bool   `json:"heartbeatMetricsEnabled,omitempty"`
	ClientMetricsEnabled    bool   `json:"clientMetricsEnabled,omitempty"`
}

type ConformanceClient struct {
	Version string `json:"version"`
}

func NewConformance(
	apiVersion string,
	infrastructureVersion uint64,
	releaseStageName string,
	serverTitle, serverName, serverArchitecture string,
	serverID uint64,
	shardsCount int,
	startedAt string,
	heartbeatMetricsEnabled,
	clientMetricsEnabled bool) (*Conformance, error) {
	if shardsCount < 0 {
		return nil, errors.New("unexpected shard servers count")
	}

	v := &Conformance{
		Name: serverTitle,
		API:  &ConformanceClient{Version: apiVersion},
		Server: &ConformanceServer{
			Type:                    serverName,
			ID:                      serverID,
			Version:                 infrastructureVersion,
			ReleaseStage:            releaseStageName,
			Sharded:                 shardsCount > 0,
			ShardsCount:             uint64(shardsCount),
			Architecture:            serverArchitecture,
			StartedAt:               startedAt,
			HeartbeatMetricsEnabled: heartbeatMetricsEnabled,
			ClientMetricsEnabled:    clientMetricsEnabled}}

	return v, nil
}
<|endoftext|>


<commit_before>package sudoku

import (
	"testing"
)

func TestCopyWithModifications(t *testing.T) {
	sourceGrid := NewGrid()
	sourceGrid.MutableCell(0, 0).SetMark(1, true)
	sourceGrid.MutableCell(0, 0).SetExcluded(1, true)

	gridFive := NewGrid()
	gridFive.MutableCell(0, 0).SetNumber(5)

	gridMarks := NewGrid()
	gridMarksCell := gridMarks.MutableCell(0, 0)
	gridMarksCell.SetMark(2, true)
	gridMarksCell.SetMark(1, false)

	gridExcludes := sourceGrid.MutableCopy()
	gridExcludesCell := gridExcludes.MutableCell(0, 0)
	gridExcludesCell.SetExcluded(2, true)
	gridExcludesCell.SetExcluded(1, false)

	tests := []struct {
		modifications GridModifcation
		expected      Grid
		description   string
	}{
		{
			GridModifcation{
				&CellModification{
					Cell:   sourceGrid.Cell(0, 0),
					Number: 5,
				},
			},
			gridFive,
			"Single valid number",
		},
		{
			GridModifcation{
				&CellModification{
					Cell:   sourceGrid.Cell(0, 0),
					Number: DIM,
				},
			},
			sourceGrid,
			"Single invalid number",
		},
		{
			GridModifcation{
				&CellModification{
					Cell:   sourceGrid.Cell(0, 0),
					Number: -1,
					MarksChanges: map[int]bool{
						1:       false,
						2:       true,
						DIM + 1: true,
					},
				},
			},
			gridMarks,
			"Marks",
		},
		{
			GridModifcation{
				&CellModification{
					Cell:   sourceGrid.Cell(0, 0),
					Number: -1,
					ExcludesChanges: map[int]bool{
						1:       false,
						2:       true,
						DIM + 1: true,
					},
				},
			},
			gridExcludes,
			"Excludes",
		},
	}

	for i, test := range tests {
		result := sourceGrid.CopyWithModifications(test.modifications)
		if result.Diagram(true) != test.expected.Diagram(true) {
			t.Error("Test", i, "failed", test.description, "Got", result.Diagram(true), "expected", test.expected.Diagram(true))
		}
	}

}
<commit_msg>Fixed one failing test<commit_after>package sudoku

import (
	"testing"
)

func TestCopyWithModifications(t *testing.T) {
	sourceGrid := NewGrid()
	sourceGrid.MutableCell(0, 0).SetMark(1, true)
	sourceGrid.MutableCell(0, 0).SetExcluded(1, true)

	gridFive := NewGrid()
	gridFive.MutableCell(0, 0).SetNumber(5)

	gridMarks := NewGrid()
	gridMarksCell := gridMarks.MutableCell(0, 0)
	gridMarksCell.SetMark(2, true)
	gridMarksCell.SetMark(1, false)

	gridExcludes := sourceGrid.MutableCopy()
	gridExcludesCell := gridExcludes.MutableCell(0, 0)
	gridExcludesCell.SetExcluded(2, true)
	gridExcludesCell.SetExcluded(1, false)

	tests := []struct {
		modifications GridModifcation
		expected      Grid
		description   string
	}{
		{
			GridModifcation{
				&CellModification{
					Cell:   sourceGrid.Cell(0, 0),
					Number: 5,
				},
			},
			gridFive,
			"Single valid number",
		},
		{
			GridModifcation{
				&CellModification{
					Cell:   sourceGrid.Cell(0, 0),
					Number: DIM + 1,
				},
			},
			sourceGrid,
			"Single invalid number",
		},
		{
			GridModifcation{
				&CellModification{
					Cell:   sourceGrid.Cell(0, 0),
					Number: -1,
					MarksChanges: map[int]bool{
						1:       false,
						2:       true,
						DIM + 1: true,
					},
				},
			},
			gridMarks,
			"Marks",
		},
		{
			GridModifcation{
				&CellModification{
					Cell:   sourceGrid.Cell(0, 0),
					Number: -1,
					ExcludesChanges: map[int]bool{
						1:       false,
						2:       true,
						DIM + 1: true,
					},
				},
			},
			gridExcludes,
			"Excludes",
		},
	}

	for i, test := range tests {
		result := sourceGrid.CopyWithModifications(test.modifications)
		if result.Diagram(true) != test.expected.Diagram(true) {
			t.Error("Test", i, "failed", test.description, "Got", result.Diagram(true), "expected", test.expected.Diagram(true))
		}
	}

}
<|endoftext|>


<commit_before>/*

Copyright 2020 The Vouch Proxy Authors.
Use of this source code is governed by The MIT License (MIT) that
can be found in the LICENSE file. Software distributed under The
MIT License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
OR CONDITIONS OF ANY KIND, either express or implied.

*/

package handlers

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

	"github.com/stretchr/testify/assert"
	vegeta "github.com/tsenart/vegeta/lib"

	"github.com/vouch/vouch-proxy/pkg/cfg"
	"github.com/vouch/vouch-proxy/pkg/jwtmanager"
	"github.com/vouch/vouch-proxy/pkg/response"
	"github.com/vouch/vouch-proxy/pkg/structs"
)

func TestValidateRequestHandlerPerf(t *testing.T) {
	setUp("/config/testing/handler_claims.yml")
	user := &structs.User{Username: "testuser", Email: "test@example.com", Name: "Test Name"}
	tokens := structs.PTokens{}
	customClaims := structs.CustomClaims{}

	userTokenString := jwtmanager.CreateUserTokenString(*user, customClaims, tokens)

	c := &http.Cookie{
		// Name:    cfg.Cfg.Cookie.Name + "_1of1",
		Name:    cfg.Cfg.Cookie.Name,
		Value:   userTokenString,
		Expires: time.Now().Add(1 * time.Hour),
	}

	// handler := http.HandlerFunc(ValidateRequestHandler)
	handler := response.JWTCacheHandler(http.HandlerFunc(ValidateRequestHandler))
	ts := httptest.NewServer(handler)
	defer ts.Close()

	rate := vegeta.Rate{Freq: 100, Per: time.Second}
	duration := 10 * time.Second
	h := &http.Header{}
	h.Add("Cookie", c.String())
	targeter := vegeta.NewStaticTargeter(vegeta.Target{
		Method: "GET",
		URL:    ts.URL,
		Header: *h,
	})

	attacker := vegeta.NewAttacker()

	var metrics vegeta.Metrics
	mustFail := false
	for res := range attacker.Attack(targeter, rate, duration, "Big Bang!") {
		if res.Code != http.StatusOK {
			t.Logf("/validate perf %d response code %d", res.Seq, res.Code)
			mustFail = true
		}
		metrics.Add(res)
	}
	metrics.Close()

	limit := time.Millisecond
	if mustFail || metrics.Latencies.P95 > limit {
		t.Logf("99th percentile latencies: %s", metrics.Latencies.P99)
		t.Logf("95th percentile latencies: %s", metrics.Latencies.P95)
		t.Logf("50th percentile latencies: %s", metrics.Latencies.P50)
		t.Logf("min latencies: %s", metrics.Latencies.Min)
		t.Logf("max latencies: %s", metrics.Latencies.Max)
		t.Logf("/validate 95th percentile latency is higher than %s", limit)
		if mustFail {
			t.Logf("not all requests were %d", http.StatusOK)
		}
		t.FailNow()
	}

}

func TestValidateRequestHandlerWithGroupClaims(t *testing.T) {
	setUp("/config/testing/handler_claims.yml")

	customClaims := structs.CustomClaims{
		Claims: map[string]interface{}{
			"sub": "f:a95afe53-60ba-4ac6-af15-fab870e72f3d:mrtester",
			"groups": []string{
				"Website Users",
				"Test Group",
			},
			"given_name":    "Mister",
			"family_name":   "Tester",
			"email":         "mrtester@test.int",
			"boolean_claim": true,
			// Auth0 custom claim are URLs
			// https://auth0.com/docs/tokens/guides/create-namespaced-custom-claims
			"http://www.example.com/favorite_color": "blue",
		},
	}

	groupHeader := "X-Vouch-IdP-Claims-Groups"
	booleanHeader := "X-Vouch-IdP-Claims-Boolean-Claim"
	familyNameHeader := "X-Vouch-IdP-Claims-Family-Name"
	favoriteColorHeader := "X-Vouch-IdP-Claims-Www-Example-Com-Favorite-Color"

	tokens := structs.PTokens{}

	user := &structs.User{Username: "testuser", Email: "test@example.com", Name: "Test Name"}
	userTokenString := jwtmanager.CreateUserTokenString(*user, customClaims, tokens)

	req, err := http.NewRequest("GET", "/validate", nil)
	if err != nil {
		t.Fatal(err)
	}

	req.AddCookie(&http.Cookie{
		// Name:    cfg.Cfg.Cookie.Name + "_1of1",
		Name:    cfg.Cfg.Cookie.Name,
		Value:   userTokenString,
		Expires: time.Now().Add(1 * time.Hour),
	})

	rr := httptest.NewRecorder()

	handler := http.HandlerFunc(ValidateRequestHandler)
	handler.ServeHTTP(rr, req)

	if status := rr.Code; status != http.StatusOK {
		t.Errorf("handler returned wrong status code: got %v want %v",
			status, http.StatusOK)
	}

	// Check that the custom claim headers are what we expected
	customClaimHeaders := map[string][]string{
		strings.ToLower(groupHeader):         {},
		strings.ToLower(booleanHeader):       {},
		strings.ToLower(familyNameHeader):    {},
		strings.ToLower(favoriteColorHeader): {},
	}

	for k, v := range rr.Result().Header {
		k = strings.ToLower(k)
		if _, exist := customClaimHeaders[k]; exist {
			customClaimHeaders[k] = v
		}
	}
	expectedCustomClaimHeaders := map[string][]string{
		strings.ToLower(groupHeader):         {"\"Website Users\",\"Test Group\""},
		strings.ToLower(booleanHeader):       {"true"},
		strings.ToLower(familyNameHeader):    {"Tester"},
		strings.ToLower(favoriteColorHeader): {"blue"},
	}
	assert.Equal(t, expectedCustomClaimHeaders, customClaimHeaders)
}
<commit_msg>add benchmark test for jwtcache'd /validate<commit_after>/*

Copyright 2020 The Vouch Proxy Authors.
Use of this source code is governed by The MIT License (MIT) that
can be found in the LICENSE file. Software distributed under The
MIT License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
OR CONDITIONS OF ANY KIND, either express or implied.

*/

package handlers

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

	"github.com/stretchr/testify/assert"
	vegeta "github.com/tsenart/vegeta/lib"

	"github.com/vouch/vouch-proxy/pkg/cfg"
	"github.com/vouch/vouch-proxy/pkg/jwtmanager"
	"github.com/vouch/vouch-proxy/pkg/structs"
)

func BenchmarkValidateRequestHandler(b *testing.B) {
	setUp("/config/testing/handler_email.yml")
	user := &structs.User{Username: "testuser", Email: "test@example.com", Name: "Test Name"}
	tokens := structs.PTokens{}
	customClaims := structs.CustomClaims{}

	userTokenString := jwtmanager.CreateUserTokenString(*user, customClaims, tokens)

	c := &http.Cookie{
		// Name:    cfg.Cfg.Cookie.Name + "_1of1",
		Name:    cfg.Cfg.Cookie.Name,
		Value:   userTokenString,
		Expires: time.Now().Add(1 * time.Hour),
	}

	handler := jwtmanager.JWTCacheHandler(http.HandlerFunc(ValidateRequestHandler))
	// handler := http.HandlerFunc(ValidateRequestHandler)
	ts := httptest.NewServer(handler)
	defer ts.Close()

	req, err := http.NewRequest("GET", "/validate", nil)
	if err != nil {
		b.Fatal(err)
	}
	req.Host = "myapp.example.com"
	req.AddCookie(c)
	w := httptest.NewRecorder()

	for i := 0; i < b.N; i++ {
		handler.ServeHTTP(w, req)
	}

}

func TestValidateRequestHandlerPerf(t *testing.T) {
	setUp("/config/testing/handler_email.yml")
	user := &structs.User{Username: "testuser", Email: "test@example.com", Name: "Test Name"}
	tokens := structs.PTokens{}
	customClaims := structs.CustomClaims{}

	userTokenString := jwtmanager.CreateUserTokenString(*user, customClaims, tokens)

	c := &http.Cookie{
		// Name:    cfg.Cfg.Cookie.Name + "_1of1",
		Name:    cfg.Cfg.Cookie.Name,
		Value:   userTokenString,
		Expires: time.Now().Add(1 * time.Hour),
	}

	// handler := http.HandlerFunc(ValidateRequestHandler)
	handler := jwtmanager.JWTCacheHandler(http.HandlerFunc(ValidateRequestHandler))
	ts := httptest.NewServer(handler)
	defer ts.Close()

	rate := vegeta.Rate{Freq: 1000, Per: time.Second}
	duration := 5 * time.Second
	h := &http.Header{}
	h.Add("Cookie", c.String())
	h.Add("Host", "myapp.example.com")
	targeter := vegeta.NewStaticTargeter(vegeta.Target{
		Method: "GET",
		URL:    ts.URL,
		Header: *h,
	})

	attacker := vegeta.NewAttacker()

	var metrics vegeta.Metrics
	mustFail := false
	for res := range attacker.Attack(targeter, rate, duration, "Big Bang!") {
		if res.Code != http.StatusOK {
			t.Logf("/validate perf %d response code %d", res.Seq, res.Code)
			mustFail = true
		}
		metrics.Add(res)
	}
	metrics.Close()

	limit := time.Millisecond
	if mustFail || metrics.Latencies.P95 > limit {
		t.Logf("99th percentile latencies: %s", metrics.Latencies.P99)
		t.Logf("95th percentile latencies: %s", metrics.Latencies.P95)
		t.Logf("50th percentile latencies: %s", metrics.Latencies.P50)
		t.Logf("min latencies: %s", metrics.Latencies.Min)
		t.Logf("max latencies: %s", metrics.Latencies.Max)
		t.Logf("/validate 95th percentile latency is higher than %s", limit)
		if mustFail {
			t.Logf("not all requests were %d", http.StatusOK)
		}
		t.FailNow()
	}

}

func TestValidateRequestHandlerWithGroupClaims(t *testing.T) {
	setUp("/config/testing/handler_claims.yml")

	customClaims := structs.CustomClaims{
		Claims: map[string]interface{}{
			"sub": "f:a95afe53-60ba-4ac6-af15-fab870e72f3d:mrtester",
			"groups": []string{
				"Website Users",
				"Test Group",
			},
			"given_name":    "Mister",
			"family_name":   "Tester",
			"email":         "mrtester@test.int",
			"boolean_claim": true,
			// Auth0 custom claim are URLs
			// https://auth0.com/docs/tokens/guides/create-namespaced-custom-claims
			"http://www.example.com/favorite_color": "blue",
		},
	}

	groupHeader := "X-Vouch-IdP-Claims-Groups"
	booleanHeader := "X-Vouch-IdP-Claims-Boolean-Claim"
	familyNameHeader := "X-Vouch-IdP-Claims-Family-Name"
	favoriteColorHeader := "X-Vouch-IdP-Claims-Www-Example-Com-Favorite-Color"

	tokens := structs.PTokens{}

	user := &structs.User{Username: "testuser", Email: "test@example.com", Name: "Test Name"}
	userTokenString := jwtmanager.CreateUserTokenString(*user, customClaims, tokens)

	req, err := http.NewRequest("GET", "/validate", nil)
	if err != nil {
		t.Fatal(err)
	}

	req.AddCookie(&http.Cookie{
		// Name:    cfg.Cfg.Cookie.Name + "_1of1",
		Name:    cfg.Cfg.Cookie.Name,
		Value:   userTokenString,
		Expires: time.Now().Add(1 * time.Hour),
	})

	rr := httptest.NewRecorder()

	handler := http.HandlerFunc(ValidateRequestHandler)
	handler.ServeHTTP(rr, req)

	if status := rr.Code; status != http.StatusOK {
		t.Errorf("handler returned wrong status code: got %v want %v",
			status, http.StatusOK)
	}

	// Check that the custom claim headers are what we expected
	customClaimHeaders := map[string][]string{
		strings.ToLower(groupHeader):         {},
		strings.ToLower(booleanHeader):       {},
		strings.ToLower(familyNameHeader):    {},
		strings.ToLower(favoriteColorHeader): {},
	}

	for k, v := range rr.Result().Header {
		k = strings.ToLower(k)
		if _, exist := customClaimHeaders[k]; exist {
			customClaimHeaders[k] = v
		}
	}
	expectedCustomClaimHeaders := map[string][]string{
		strings.ToLower(groupHeader):         {"\"Website Users\",\"Test Group\""},
		strings.ToLower(booleanHeader):       {"true"},
		strings.ToLower(familyNameHeader):    {"Tester"},
		strings.ToLower(favoriteColorHeader): {"blue"},
	}
	assert.Equal(t, expectedCustomClaimHeaders, customClaimHeaders)
}
<|endoftext|>


<commit_before>package handshake

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"io"
	"net"
	"sync"

	"github.com/lucas-clemente/quic-go/crypto"
	"github.com/lucas-clemente/quic-go/protocol"
	"github.com/lucas-clemente/quic-go/qerr"
	"github.com/lucas-clemente/quic-go/utils"
)

// KeyDerivationFunction is used for key derivation
type KeyDerivationFunction func(forwardSecure bool, sharedSecret, nonces []byte, connID protocol.ConnectionID, chlo []byte, scfg []byte, cert []byte, divNonce []byte) (crypto.AEAD, error)

// KeyExchangeFunction is used to make a new KEX
type KeyExchangeFunction func() crypto.KeyExchange

// The CryptoSetup handles all things crypto for the Session
type CryptoSetup struct {
	connID               protocol.ConnectionID
	ip                   net.IP
	version              protocol.VersionNumber
	scfg                 *ServerConfig
	diversificationNonce []byte

	secureAEAD                  crypto.AEAD
	forwardSecureAEAD           crypto.AEAD
	receivedForwardSecurePacket bool
	receivedSecurePacket        bool
	aeadChanged                 chan struct{}

	keyDerivation KeyDerivationFunction
	keyExchange   KeyExchangeFunction

	cryptoStream utils.Stream

	connectionParameters ConnectionParametersManager

	mutex sync.RWMutex
}

var _ crypto.AEAD = &CryptoSetup{}

// NewCryptoSetup creates a new CryptoSetup instance
func NewCryptoSetup(
	connID protocol.ConnectionID,
	ip net.IP,
	version protocol.VersionNumber,
	scfg *ServerConfig,
	cryptoStream utils.Stream,
	connectionParameters ConnectionParametersManager,
	aeadChanged chan struct{},
) (*CryptoSetup, error) {
	return &CryptoSetup{
		connID:               connID,
		ip:                   ip,
		version:              version,
		scfg:                 scfg,
		keyDerivation:        crypto.DeriveKeysAESGCM,
		keyExchange:          getEphermalKEX,
		cryptoStream:         cryptoStream,
		connectionParameters: connectionParameters,
		aeadChanged:          aeadChanged,
	}, nil
}

// HandleCryptoStream reads and writes messages on the crypto stream
func (h *CryptoSetup) HandleCryptoStream() error {
	for {
		var chloData bytes.Buffer
		messageTag, cryptoData, err := ParseHandshakeMessage(io.TeeReader(h.cryptoStream, &chloData))
		if err != nil {
			return qerr.HandshakeFailed
		}
		if messageTag != TagCHLO {
			return qerr.InvalidCryptoMessageType
		}

		utils.Debugf("Got CHLO:\n%s", printHandshakeMessage(cryptoData))

		done, err := h.handleMessage(chloData.Bytes(), cryptoData)
		if err != nil {
			return err
		}
		if done {
			return nil
		}
	}
}

func (h *CryptoSetup) handleMessage(chloData []byte, cryptoData map[Tag][]byte) (bool, error) {
	sniSlice, ok := cryptoData[TagSNI]
	if !ok {
		return false, qerr.Error(qerr.CryptoMessageParameterNotFound, "SNI required")
	}
	sni := string(sniSlice)
	if sni == "" {
		return false, qerr.Error(qerr.CryptoMessageParameterNotFound, "SNI required")
	}

	// prevent version downgrade attacks
	// see https://groups.google.com/a/chromium.org/forum/#!topic/proto-quic/N-de9j63tCk for a discussion and examples
	verSlice, ok := cryptoData[TagVER]
	if !ok {
		return false, qerr.Error(qerr.InvalidCryptoMessageParameter, "client hello missing version tag")
	}
	if len(verSlice) != 4 {
		return false, qerr.Error(qerr.InvalidCryptoMessageParameter, "incorrect version tag")
	}
	verTag := binary.LittleEndian.Uint32(verSlice)
	ver := protocol.VersionTagToNumber(verTag)
	// If the client's preferred version is not the version we are currently speaking, then the client went through a version negotiation.  In this case, we need to make sure that we actually do not support this version and that it wasn't a downgrade attack.
	if ver != h.version && protocol.IsSupportedVersion(ver) {
		return false, qerr.Error(qerr.VersionNegotiationMismatch, "Downgrade attack detected")
	}

	var reply []byte
	var err error

	certUncompressed, err := h.scfg.signer.GetLeafCert(sni)
	if err != nil {
		return false, err
	}

	if !h.isInchoateCHLO(cryptoData, certUncompressed) {
		// We have a CHLO with a proper server config ID, do a 0-RTT handshake
		reply, err = h.handleCHLO(sni, chloData, cryptoData)
		if err != nil {
			return false, err
		}
		_, err = h.cryptoStream.Write(reply)
		if err != nil {
			return false, err
		}
		return true, nil
	}

	// We have an inchoate or non-matching CHLO, we now send a rejection
	reply, err = h.handleInchoateCHLO(sni, chloData, cryptoData)
	if err != nil {
		return false, err
	}
	_, err = h.cryptoStream.Write(reply)
	if err != nil {
		return false, err
	}
	return false, nil
}

// Open a message
func (h *CryptoSetup) Open(dst, src []byte, packetNumber protocol.PacketNumber, associatedData []byte) ([]byte, error) {
	h.mutex.RLock()
	defer h.mutex.RUnlock()

	if h.forwardSecureAEAD != nil {
		res, err := h.forwardSecureAEAD.Open(dst, src, packetNumber, associatedData)
		if err == nil {
			h.receivedForwardSecurePacket = true
			return res, nil
		}
		if h.receivedForwardSecurePacket {
			return nil, err
		}
	}
	if h.secureAEAD != nil {
		res, err := h.secureAEAD.Open(dst, src, packetNumber, associatedData)
		if err == nil {
			h.receivedSecurePacket = true
			return res, nil
		}
		if h.receivedSecurePacket {
			return nil, err
		}
	}
	return (&crypto.NullAEAD{}).Open(dst, src, packetNumber, associatedData)
}

// Seal a message, call LockForSealing() before!
func (h *CryptoSetup) Seal(dst, src []byte, packetNumber protocol.PacketNumber, associatedData []byte) []byte {
	if h.receivedForwardSecurePacket {
		return h.forwardSecureAEAD.Seal(dst, src, packetNumber, associatedData)
	} else if h.secureAEAD != nil {
		return h.secureAEAD.Seal(dst, src, packetNumber, associatedData)
	} else {
		return (&crypto.NullAEAD{}).Seal(dst, src, packetNumber, associatedData)
	}
}

func (h *CryptoSetup) isInchoateCHLO(cryptoData map[Tag][]byte, cert []byte) bool {
	if _, ok := cryptoData[TagPUBS]; !ok {
		return true
	}
	scid, ok := cryptoData[TagSCID]
	if !ok || !bytes.Equal(h.scfg.ID, scid) {
		return true
	}
	xlctTag, ok := cryptoData[TagXLCT]
	if !ok || len(xlctTag) != 8 {
		return true
	}
	xlct := binary.LittleEndian.Uint64(xlctTag)
	if crypto.HashCert(cert) != xlct {
		return true
	}
	if err := h.scfg.stkSource.VerifyToken(h.ip, cryptoData[TagSTK]); err != nil {
		utils.Infof("STK invalid: %s", err.Error())
		return true
	}
	return false
}

func (h *CryptoSetup) handleInchoateCHLO(sni string, chlo []byte, cryptoData map[Tag][]byte) ([]byte, error) {
	if len(chlo) < protocol.ClientHelloMinimumSize {
		return nil, qerr.Error(qerr.CryptoInvalidValueLength, "CHLO too small")
	}

	token, err := h.scfg.stkSource.NewToken(h.ip)
	if err != nil {
		return nil, err
	}

	replyMap := map[Tag][]byte{
		TagSCFG: h.scfg.Get(),
		TagSTK:  token,
		TagSVID: []byte("quic-go"),
	}

	if h.scfg.stkSource.VerifyToken(h.ip, cryptoData[TagSTK]) == nil {
		proof, err := h.scfg.Sign(sni, chlo)
		if err != nil {
			return nil, err
		}

		commonSetHashes := cryptoData[TagCCS]
		cachedCertsHashes := cryptoData[TagCCRT]

		certCompressed, err := h.scfg.GetCertsCompressed(sni, commonSetHashes, cachedCertsHashes)
		if err != nil {
			return nil, err
		}
		// Token was valid, send more details
		replyMap[TagPROF] = proof
		replyMap[TagCERT] = certCompressed
	}

	var serverReply bytes.Buffer
	WriteHandshakeMessage(&serverReply, TagREJ, replyMap)
	return serverReply.Bytes(), nil
}

func (h *CryptoSetup) handleCHLO(sni string, data []byte, cryptoData map[Tag][]byte) ([]byte, error) {
	// We have a CHLO matching our server config, we can continue with the 0-RTT handshake
	sharedSecret, err := h.scfg.kex.CalculateSharedKey(cryptoData[TagPUBS])
	if err != nil {
		return nil, err
	}

	h.mutex.Lock()
	defer h.mutex.Unlock()

	certUncompressed, err := h.scfg.signer.GetLeafCert(sni)
	if err != nil {
		return nil, err
	}

	serverNonce := make([]byte, 32)
	if _, err = rand.Read(serverNonce); err != nil {
		return nil, err
	}

	h.diversificationNonce = make([]byte, 32)
	if _, err = rand.Read(h.diversificationNonce); err != nil {
		return nil, err
	}

	clientNonce := cryptoData[TagNONC]
	err = h.validateClientNonce(clientNonce)
	if err != nil {
		return nil, err
	}

	aead := cryptoData[TagAEAD]
	if !bytes.Equal(aead, []byte("AESG")) {
		return nil, qerr.Error(qerr.CryptoNoSupport, "Unsupported AEAD or KEXS")
	}

	kexs := cryptoData[TagKEXS]
	if !bytes.Equal(kexs, []byte("C255")) {
		return nil, qerr.Error(qerr.CryptoNoSupport, "Unsupported AEAD or KEXS")
	}

	h.secureAEAD, err = h.keyDerivation(
		false,
		sharedSecret,
		clientNonce,
		h.connID,
		data,
		h.scfg.Get(),
		certUncompressed,
		h.diversificationNonce,
	)
	if err != nil {
		return nil, err
	}

	// Generate a new curve instance to derive the forward secure key
	var fsNonce bytes.Buffer
	fsNonce.Write(clientNonce)
	fsNonce.Write(serverNonce)
	ephermalKex := h.keyExchange()
	ephermalSharedSecret, err := ephermalKex.CalculateSharedKey(cryptoData[TagPUBS])
	if err != nil {
		return nil, err
	}

	h.forwardSecureAEAD, err = h.keyDerivation(
		true,
		ephermalSharedSecret,
		fsNonce.Bytes(),
		h.connID,
		data,
		h.scfg.Get(),
		certUncompressed,
		nil,
	)
	if err != nil {
		return nil, err
	}

	err = h.connectionParameters.SetFromMap(cryptoData)
	if err != nil {
		return nil, err
	}

	replyMap := h.connectionParameters.GetSHLOMap()
	// add crypto parameters
	replyMap[TagPUBS] = ephermalKex.PublicKey()
	replyMap[TagSNO] = serverNonce
	replyMap[TagVER] = protocol.SupportedVersionsAsTags

	var reply bytes.Buffer
	WriteHandshakeMessage(&reply, TagSHLO, replyMap)

	h.aeadChanged <- struct{}{}

	return reply.Bytes(), nil
}

// DiversificationNonce returns a diversification nonce if required in the next packet to be Seal'ed. See LockForSealing()!
func (h *CryptoSetup) DiversificationNonce() []byte {
	if h.receivedForwardSecurePacket || h.secureAEAD == nil {
		return nil
	}
	return h.diversificationNonce
}

// LockForSealing should be called before Seal(). It is needed so that diversification nonces can be obtained before packets are sealed, and the AEADs are not changed in the meantime.
func (h *CryptoSetup) LockForSealing() {
	h.mutex.RLock()
}

// UnlockForSealing should be called after Seal() is complete, see LockForSealing().
func (h *CryptoSetup) UnlockForSealing() {
	h.mutex.RUnlock()
}

// HandshakeComplete returns true after the first forward secure packet was received form the client.
func (h *CryptoSetup) HandshakeComplete() bool {
	return h.receivedForwardSecurePacket
}

func (h *CryptoSetup) validateClientNonce(nonce []byte) error {
	if len(nonce) != 32 {
		return qerr.Error(qerr.InvalidCryptoMessageParameter, "invalid client nonce length")
	}
	if !bytes.Equal(nonce[4:12], h.scfg.obit) {
		return qerr.Error(qerr.InvalidCryptoMessageParameter, "OBIT not matching")
	}
	return nil
}
<commit_msg>log REJs and SHLOs<commit_after>package handshake

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"io"
	"net"
	"sync"

	"github.com/lucas-clemente/quic-go/crypto"
	"github.com/lucas-clemente/quic-go/protocol"
	"github.com/lucas-clemente/quic-go/qerr"
	"github.com/lucas-clemente/quic-go/utils"
)

// KeyDerivationFunction is used for key derivation
type KeyDerivationFunction func(forwardSecure bool, sharedSecret, nonces []byte, connID protocol.ConnectionID, chlo []byte, scfg []byte, cert []byte, divNonce []byte) (crypto.AEAD, error)

// KeyExchangeFunction is used to make a new KEX
type KeyExchangeFunction func() crypto.KeyExchange

// The CryptoSetup handles all things crypto for the Session
type CryptoSetup struct {
	connID               protocol.ConnectionID
	ip                   net.IP
	version              protocol.VersionNumber
	scfg                 *ServerConfig
	diversificationNonce []byte

	secureAEAD                  crypto.AEAD
	forwardSecureAEAD           crypto.AEAD
	receivedForwardSecurePacket bool
	receivedSecurePacket        bool
	aeadChanged                 chan struct{}

	keyDerivation KeyDerivationFunction
	keyExchange   KeyExchangeFunction

	cryptoStream utils.Stream

	connectionParameters ConnectionParametersManager

	mutex sync.RWMutex
}

var _ crypto.AEAD = &CryptoSetup{}

// NewCryptoSetup creates a new CryptoSetup instance
func NewCryptoSetup(
	connID protocol.ConnectionID,
	ip net.IP,
	version protocol.VersionNumber,
	scfg *ServerConfig,
	cryptoStream utils.Stream,
	connectionParameters ConnectionParametersManager,
	aeadChanged chan struct{},
) (*CryptoSetup, error) {
	return &CryptoSetup{
		connID:               connID,
		ip:                   ip,
		version:              version,
		scfg:                 scfg,
		keyDerivation:        crypto.DeriveKeysAESGCM,
		keyExchange:          getEphermalKEX,
		cryptoStream:         cryptoStream,
		connectionParameters: connectionParameters,
		aeadChanged:          aeadChanged,
	}, nil
}

// HandleCryptoStream reads and writes messages on the crypto stream
func (h *CryptoSetup) HandleCryptoStream() error {
	for {
		var chloData bytes.Buffer
		messageTag, cryptoData, err := ParseHandshakeMessage(io.TeeReader(h.cryptoStream, &chloData))
		if err != nil {
			return qerr.HandshakeFailed
		}
		if messageTag != TagCHLO {
			return qerr.InvalidCryptoMessageType
		}

		utils.Debugf("Got CHLO:\n%s", printHandshakeMessage(cryptoData))

		done, err := h.handleMessage(chloData.Bytes(), cryptoData)
		if err != nil {
			return err
		}
		if done {
			return nil
		}
	}
}

func (h *CryptoSetup) handleMessage(chloData []byte, cryptoData map[Tag][]byte) (bool, error) {
	sniSlice, ok := cryptoData[TagSNI]
	if !ok {
		return false, qerr.Error(qerr.CryptoMessageParameterNotFound, "SNI required")
	}
	sni := string(sniSlice)
	if sni == "" {
		return false, qerr.Error(qerr.CryptoMessageParameterNotFound, "SNI required")
	}

	// prevent version downgrade attacks
	// see https://groups.google.com/a/chromium.org/forum/#!topic/proto-quic/N-de9j63tCk for a discussion and examples
	verSlice, ok := cryptoData[TagVER]
	if !ok {
		return false, qerr.Error(qerr.InvalidCryptoMessageParameter, "client hello missing version tag")
	}
	if len(verSlice) != 4 {
		return false, qerr.Error(qerr.InvalidCryptoMessageParameter, "incorrect version tag")
	}
	verTag := binary.LittleEndian.Uint32(verSlice)
	ver := protocol.VersionTagToNumber(verTag)
	// If the client's preferred version is not the version we are currently speaking, then the client went through a version negotiation.  In this case, we need to make sure that we actually do not support this version and that it wasn't a downgrade attack.
	if ver != h.version && protocol.IsSupportedVersion(ver) {
		return false, qerr.Error(qerr.VersionNegotiationMismatch, "Downgrade attack detected")
	}

	var reply []byte
	var err error

	certUncompressed, err := h.scfg.signer.GetLeafCert(sni)
	if err != nil {
		return false, err
	}

	if !h.isInchoateCHLO(cryptoData, certUncompressed) {
		// We have a CHLO with a proper server config ID, do a 0-RTT handshake
		reply, err = h.handleCHLO(sni, chloData, cryptoData)
		if err != nil {
			return false, err
		}
		_, err = h.cryptoStream.Write(reply)
		if err != nil {
			return false, err
		}
		return true, nil
	}

	// We have an inchoate or non-matching CHLO, we now send a rejection
	reply, err = h.handleInchoateCHLO(sni, chloData, cryptoData)
	if err != nil {
		return false, err
	}
	_, err = h.cryptoStream.Write(reply)
	if err != nil {
		return false, err
	}
	return false, nil
}

// Open a message
func (h *CryptoSetup) Open(dst, src []byte, packetNumber protocol.PacketNumber, associatedData []byte) ([]byte, error) {
	h.mutex.RLock()
	defer h.mutex.RUnlock()

	if h.forwardSecureAEAD != nil {
		res, err := h.forwardSecureAEAD.Open(dst, src, packetNumber, associatedData)
		if err == nil {
			h.receivedForwardSecurePacket = true
			return res, nil
		}
		if h.receivedForwardSecurePacket {
			return nil, err
		}
	}
	if h.secureAEAD != nil {
		res, err := h.secureAEAD.Open(dst, src, packetNumber, associatedData)
		if err == nil {
			h.receivedSecurePacket = true
			return res, nil
		}
		if h.receivedSecurePacket {
			return nil, err
		}
	}
	return (&crypto.NullAEAD{}).Open(dst, src, packetNumber, associatedData)
}

// Seal a message, call LockForSealing() before!
func (h *CryptoSetup) Seal(dst, src []byte, packetNumber protocol.PacketNumber, associatedData []byte) []byte {
	if h.receivedForwardSecurePacket {
		return h.forwardSecureAEAD.Seal(dst, src, packetNumber, associatedData)
	} else if h.secureAEAD != nil {
		return h.secureAEAD.Seal(dst, src, packetNumber, associatedData)
	} else {
		return (&crypto.NullAEAD{}).Seal(dst, src, packetNumber, associatedData)
	}
}

func (h *CryptoSetup) isInchoateCHLO(cryptoData map[Tag][]byte, cert []byte) bool {
	if _, ok := cryptoData[TagPUBS]; !ok {
		return true
	}
	scid, ok := cryptoData[TagSCID]
	if !ok || !bytes.Equal(h.scfg.ID, scid) {
		return true
	}
	xlctTag, ok := cryptoData[TagXLCT]
	if !ok || len(xlctTag) != 8 {
		return true
	}
	xlct := binary.LittleEndian.Uint64(xlctTag)
	if crypto.HashCert(cert) != xlct {
		return true
	}
	if err := h.scfg.stkSource.VerifyToken(h.ip, cryptoData[TagSTK]); err != nil {
		utils.Infof("STK invalid: %s", err.Error())
		return true
	}
	return false
}

func (h *CryptoSetup) handleInchoateCHLO(sni string, chlo []byte, cryptoData map[Tag][]byte) ([]byte, error) {
	if len(chlo) < protocol.ClientHelloMinimumSize {
		return nil, qerr.Error(qerr.CryptoInvalidValueLength, "CHLO too small")
	}

	token, err := h.scfg.stkSource.NewToken(h.ip)
	if err != nil {
		return nil, err
	}

	replyMap := map[Tag][]byte{
		TagSCFG: h.scfg.Get(),
		TagSTK:  token,
		TagSVID: []byte("quic-go"),
	}

	if h.scfg.stkSource.VerifyToken(h.ip, cryptoData[TagSTK]) == nil {
		proof, err := h.scfg.Sign(sni, chlo)
		if err != nil {
			return nil, err
		}

		commonSetHashes := cryptoData[TagCCS]
		cachedCertsHashes := cryptoData[TagCCRT]

		certCompressed, err := h.scfg.GetCertsCompressed(sni, commonSetHashes, cachedCertsHashes)
		if err != nil {
			return nil, err
		}
		// Token was valid, send more details
		replyMap[TagPROF] = proof
		replyMap[TagCERT] = certCompressed
	}

	var serverReply bytes.Buffer
	WriteHandshakeMessage(&serverReply, TagREJ, replyMap)
	utils.Debugf("Sending REJ:\n%s", printHandshakeMessage(cryptoData))
	return serverReply.Bytes(), nil
}

func (h *CryptoSetup) handleCHLO(sni string, data []byte, cryptoData map[Tag][]byte) ([]byte, error) {
	// We have a CHLO matching our server config, we can continue with the 0-RTT handshake
	sharedSecret, err := h.scfg.kex.CalculateSharedKey(cryptoData[TagPUBS])
	if err != nil {
		return nil, err
	}

	h.mutex.Lock()
	defer h.mutex.Unlock()

	certUncompressed, err := h.scfg.signer.GetLeafCert(sni)
	if err != nil {
		return nil, err
	}

	serverNonce := make([]byte, 32)
	if _, err = rand.Read(serverNonce); err != nil {
		return nil, err
	}

	h.diversificationNonce = make([]byte, 32)
	if _, err = rand.Read(h.diversificationNonce); err != nil {
		return nil, err
	}

	clientNonce := cryptoData[TagNONC]
	err = h.validateClientNonce(clientNonce)
	if err != nil {
		return nil, err
	}

	aead := cryptoData[TagAEAD]
	if !bytes.Equal(aead, []byte("AESG")) {
		return nil, qerr.Error(qerr.CryptoNoSupport, "Unsupported AEAD or KEXS")
	}

	kexs := cryptoData[TagKEXS]
	if !bytes.Equal(kexs, []byte("C255")) {
		return nil, qerr.Error(qerr.CryptoNoSupport, "Unsupported AEAD or KEXS")
	}

	h.secureAEAD, err = h.keyDerivation(
		false,
		sharedSecret,
		clientNonce,
		h.connID,
		data,
		h.scfg.Get(),
		certUncompressed,
		h.diversificationNonce,
	)
	if err != nil {
		return nil, err
	}

	// Generate a new curve instance to derive the forward secure key
	var fsNonce bytes.Buffer
	fsNonce.Write(clientNonce)
	fsNonce.Write(serverNonce)
	ephermalKex := h.keyExchange()
	ephermalSharedSecret, err := ephermalKex.CalculateSharedKey(cryptoData[TagPUBS])
	if err != nil {
		return nil, err
	}

	h.forwardSecureAEAD, err = h.keyDerivation(
		true,
		ephermalSharedSecret,
		fsNonce.Bytes(),
		h.connID,
		data,
		h.scfg.Get(),
		certUncompressed,
		nil,
	)
	if err != nil {
		return nil, err
	}

	err = h.connectionParameters.SetFromMap(cryptoData)
	if err != nil {
		return nil, err
	}

	replyMap := h.connectionParameters.GetSHLOMap()
	// add crypto parameters
	replyMap[TagPUBS] = ephermalKex.PublicKey()
	replyMap[TagSNO] = serverNonce
	replyMap[TagVER] = protocol.SupportedVersionsAsTags

	var reply bytes.Buffer
	WriteHandshakeMessage(&reply, TagSHLO, replyMap)
	utils.Debugf("Sending SHLO:\n%s", printHandshakeMessage(cryptoData))

	h.aeadChanged <- struct{}{}

	return reply.Bytes(), nil
}

// DiversificationNonce returns a diversification nonce if required in the next packet to be Seal'ed. See LockForSealing()!
func (h *CryptoSetup) DiversificationNonce() []byte {
	if h.receivedForwardSecurePacket || h.secureAEAD == nil {
		return nil
	}
	return h.diversificationNonce
}

// LockForSealing should be called before Seal(). It is needed so that diversification nonces can be obtained before packets are sealed, and the AEADs are not changed in the meantime.
func (h *CryptoSetup) LockForSealing() {
	h.mutex.RLock()
}

// UnlockForSealing should be called after Seal() is complete, see LockForSealing().
func (h *CryptoSetup) UnlockForSealing() {
	h.mutex.RUnlock()
}

// HandshakeComplete returns true after the first forward secure packet was received form the client.
func (h *CryptoSetup) HandshakeComplete() bool {
	return h.receivedForwardSecurePacket
}

func (h *CryptoSetup) validateClientNonce(nonce []byte) error {
	if len(nonce) != 32 {
		return qerr.Error(qerr.InvalidCryptoMessageParameter, "invalid client nonce length")
	}
	if !bytes.Equal(nonce[4:12], h.scfg.obit) {
		return qerr.Error(qerr.InvalidCryptoMessageParameter, "OBIT not matching")
	}
	return nil
}
<|endoftext|>