<commit_before>package testutils

import (
	"net"
	"sync"
	"time"
)

// MockConn implements a net.Conn that can be used for testing.
// Deadline methods are no-ops.
type MockConn struct {
	// Function to call when Read is called.
	ReadFunc func(i int, b []byte) (int, error)
	// Function to call when Write is called.
	WriteFunc func(i int, b []byte) (int, error)

	// Error to return when Close is called.
	CloseErr error
	// If set, the channel is closed when Close is called.
	CloseChan chan struct{}

	// Local address to return when LocalAddr is called.
	LocalAddress net.Addr
	// Remote address to return when RemoteAddr is called.
	RemoteAddress net.Addr

	mu         sync.Mutex // protects close(CloseChan) and the indices
	readIndex  int
	writeIndex int
	closeIndex int
	closedAt   time.Time
}

// CloseCalls returns the number of times Close was called.
func (c *MockConn) CloseCalls() int {
	c.mu.Lock()
	i := c.closeIndex
	c.mu.Unlock()
	return i
}

func (c *MockConn) ClosedAt() time.Time {
	c.mu.Lock()
	t := c.closedAt
	c.mu.Unlock()
	return t
}

// ReadCalls returns the number of times Read was called.
func (c *MockConn) ReadCalls() int {
	c.mu.Lock()
	i := c.readIndex
	c.mu.Unlock()
	return i
}

// WriteCalls returns the number of times Write was called.
func (c *MockConn) WriteCalls() int {
	c.mu.Lock()
	i := c.writeIndex
	c.mu.Unlock()
	return i
}

// Read implements io.Reader for MockConn.
func (c *MockConn) Read(b []byte) (int, error) {
	c.mu.Lock()
	i := c.readIndex
	c.readIndex++
	c.mu.Unlock()

	if c.ReadFunc == nil {
		return 0, nil
	}
	return c.ReadFunc(i, b)
}

// Write implements io.Writer for MockConn.
func (c *MockConn) Write(b []byte) (int, error) {
	c.mu.Lock()
	i := c.writeIndex
	c.writeIndex++
	c.mu.Unlock()

	if c.WriteFunc == nil {
		return 0, nil
	}
	return c.WriteFunc(i, b)
}

// Close implements io.Closer for MockConn.
func (c *MockConn) Close() error {
	c.mu.Lock()
	if c.CloseChan != nil {
		select {
		case <-c.CloseChan:
			// already closed
		default:
			close(c.CloseChan)
			c.closedAt = time.Now()
		}
	}
	c.closeIndex++
	c.mu.Unlock()
	return c.CloseErr
}

// LocalAddr returns the local address of the connection.
func (c *MockConn) LocalAddr() net.Addr {
	return c.LocalAddress
}

// RemoteAddr returns the remote address of the connection.
func (c *MockConn) RemoteAddr() net.Addr {
	return c.RemoteAddress
}

// SetDeadline is a no-op for MockConn.
func (c *MockConn) SetDeadline(t time.Time) error {
	return nil
}

// SetReadDeadline is a no-op for MockConn.
func (c *MockConn) SetReadDeadline(t time.Time) error {
	return nil
}

// SetWriteDeadline is a no-op for MockConn.
func (c *MockConn) SetWriteDeadline(t time.Time) error {
	return nil
}
<commit_msg>testutils: add comment to ClosedAt<commit_after>package testutils

import (
	"net"
	"sync"
	"time"
)

// MockConn implements a net.Conn that can be used for testing.
// Deadline methods are no-ops.
type MockConn struct {
	// Function to call when Read is called.
	ReadFunc func(i int, b []byte) (int, error)
	// Function to call when Write is called.
	WriteFunc func(i int, b []byte) (int, error)

	// Error to return when Close is called.
	CloseErr error
	// If set, the channel is closed when Close is called.
	CloseChan chan struct{}

	// Local address to return when LocalAddr is called.
	LocalAddress net.Addr
	// Remote address to return when RemoteAddr is called.
	RemoteAddress net.Addr

	mu         sync.Mutex // protects close(CloseChan) and the indices
	readIndex  int
	writeIndex int
	closeIndex int
	closedAt   time.Time
}

// CloseCalls returns the number of times Close was called.
func (c *MockConn) CloseCalls() int {
	c.mu.Lock()
	i := c.closeIndex
	c.mu.Unlock()
	return i
}

// ClosedAt returns the timestamp at which the connection was closed.
func (c *MockConn) ClosedAt() time.Time {
	c.mu.Lock()
	t := c.closedAt
	c.mu.Unlock()
	return t
}

// ReadCalls returns the number of times Read was called.
func (c *MockConn) ReadCalls() int {
	c.mu.Lock()
	i := c.readIndex
	c.mu.Unlock()
	return i
}

// WriteCalls returns the number of times Write was called.
func (c *MockConn) WriteCalls() int {
	c.mu.Lock()
	i := c.writeIndex
	c.mu.Unlock()
	return i
}

// Read implements io.Reader for MockConn.
func (c *MockConn) Read(b []byte) (int, error) {
	c.mu.Lock()
	i := c.readIndex
	c.readIndex++
	c.mu.Unlock()

	if c.ReadFunc == nil {
		return 0, nil
	}
	return c.ReadFunc(i, b)
}

// Write implements io.Writer for MockConn.
func (c *MockConn) Write(b []byte) (int, error) {
	c.mu.Lock()
	i := c.writeIndex
	c.writeIndex++
	c.mu.Unlock()

	if c.WriteFunc == nil {
		return 0, nil
	}
	return c.WriteFunc(i, b)
}

// Close implements io.Closer for MockConn.
func (c *MockConn) Close() error {
	c.mu.Lock()
	if c.CloseChan != nil {
		select {
		case <-c.CloseChan:
			// already closed
		default:
			close(c.CloseChan)
			c.closedAt = time.Now()
		}
	}
	c.closeIndex++
	c.mu.Unlock()
	return c.CloseErr
}

// LocalAddr returns the local address of the connection.
func (c *MockConn) LocalAddr() net.Addr {
	return c.LocalAddress
}

// RemoteAddr returns the remote address of the connection.
func (c *MockConn) RemoteAddr() net.Addr {
	return c.RemoteAddress
}

// SetDeadline is a no-op for MockConn.
func (c *MockConn) SetDeadline(t time.Time) error {
	return nil
}

// SetReadDeadline is a no-op for MockConn.
func (c *MockConn) SetReadDeadline(t time.Time) error {
	return nil
}

// SetWriteDeadline is a no-op for MockConn.
func (c *MockConn) SetWriteDeadline(t time.Time) error {
	return nil
}
<|endoftext|>


<commit_before>package backup

import (
	"context"
	"io"
	"sync"
	"time"

	"github.com/restic/restic/internal/archiver"
	"github.com/restic/restic/internal/restic"
	"github.com/restic/restic/internal/ui/signals"
)

type ProgressPrinter interface {
	Update(total, processed Counter, errors uint, currentFiles map[string]struct{}, start time.Time, secs uint64)
	Error(item string, err error) error
	ScannerError(item string, err error) error
	CompleteItem(messageType string, item string, previous, current *restic.Node, s archiver.ItemStats, d time.Duration)
	ReportTotal(item string, start time.Time, s archiver.ScanStats)
	Finish(snapshotID restic.ID, start time.Time, summary *Summary, dryRun bool)
	Reset()

	// ui.StdioWrapper
	Stdout() io.WriteCloser
	Stderr() io.WriteCloser

	E(msg string, args ...interface{})
	P(msg string, args ...interface{})
	V(msg string, args ...interface{})
	VV(msg string, args ...interface{})
}

type Counter struct {
	Files, Dirs, Bytes uint64
}

type fileWorkerMessage struct {
	filename string
	done     bool
}

type ProgressReporter interface {
	CompleteItem(item string, previous, current *restic.Node, s archiver.ItemStats, d time.Duration)
	StartFile(filename string)
	CompleteBlob(filename string, bytes uint64)
	ScannerError(item string, err error) error
	ReportTotal(item string, s archiver.ScanStats)
	SetMinUpdatePause(d time.Duration)
	Run(ctx context.Context) error
	Error(item string, err error) error
	Finish(snapshotID restic.ID)
}

type Summary struct {
	sync.Mutex
	Files, Dirs struct {
		New       uint
		Changed   uint
		Unchanged uint
	}
	ProcessedBytes uint64
	archiver.ItemStats
}

// Progress reports progress for the `backup` command.
type Progress struct {
	MinUpdatePause time.Duration

	start time.Time
	dry   bool

	totalBytes uint64

	totalCh     chan Counter
	processedCh chan Counter
	errCh       chan struct{}
	workerCh    chan fileWorkerMessage
	closed      chan struct{}

	summary *Summary
	printer ProgressPrinter
}

func NewProgress(printer ProgressPrinter) *Progress {
	return &Progress{
		// limit to 60fps by default
		MinUpdatePause: time.Second / 60,
		start:          time.Now(),

		totalCh:     make(chan Counter),
		processedCh: make(chan Counter),
		errCh:       make(chan struct{}),
		workerCh:    make(chan fileWorkerMessage),
		closed:      make(chan struct{}),

		summary: &Summary{},

		printer: printer,
	}
}

// Run regularly updates the status lines. It should be called in a separate
// goroutine.
func (p *Progress) Run(ctx context.Context) error {
	var (
		lastUpdate       time.Time
		total, processed Counter
		errors           uint
		started          bool
		currentFiles     = make(map[string]struct{})
		secondsRemaining uint64
	)

	t := time.NewTicker(time.Second)
	signalsCh := signals.GetProgressChannel()
	defer t.Stop()
	defer close(p.closed)
	// Reset status when finished
	defer p.printer.Reset()

	for {
		forceUpdate := false
		select {
		case <-ctx.Done():
			return nil
		case t, ok := <-p.totalCh:
			if ok {
				total = t
				started = true
			} else {
				// scan has finished
				p.totalCh = nil
				p.totalBytes = total.Bytes
			}
		case s := <-p.processedCh:
			processed.Files += s.Files
			processed.Dirs += s.Dirs
			processed.Bytes += s.Bytes
			started = true
		case <-p.errCh:
			errors++
			started = true
		case m := <-p.workerCh:
			if m.done {
				delete(currentFiles, m.filename)
			} else {
				currentFiles[m.filename] = struct{}{}
			}
		case <-t.C:
			if !started {
				continue
			}

			if p.totalCh == nil {
				secs := float64(time.Since(p.start) / time.Second)
				todo := float64(total.Bytes - processed.Bytes)
				secondsRemaining = uint64(secs / float64(processed.Bytes) * todo)
			}
		case <-signalsCh:
			forceUpdate = true
		}

		// limit update frequency
		if !forceUpdate && (time.Since(lastUpdate) < p.MinUpdatePause || p.MinUpdatePause == 0) {
			continue
		}
		lastUpdate = time.Now()

		p.printer.Update(total, processed, errors, currentFiles, p.start, secondsRemaining)
	}
}

// ScannerError is the error callback function for the scanner, it prints the
// error in verbose mode and returns nil.
func (p *Progress) ScannerError(item string, err error) error {
	return p.printer.ScannerError(item, err)
}

// Error is the error callback function for the archiver, it prints the error and returns nil.
func (p *Progress) Error(item string, err error) error {
	cbErr := p.printer.Error(item, err)

	select {
	case p.errCh <- struct{}{}:
	case <-p.closed:
	}
	return cbErr
}

// StartFile is called when a file is being processed by a worker.
func (p *Progress) StartFile(filename string) {
	select {
	case p.workerCh <- fileWorkerMessage{filename: filename}:
	case <-p.closed:
	}
}

// CompleteBlob is called for all saved blobs for files.
func (p *Progress) CompleteBlob(filename string, bytes uint64) {
	select {
	case p.processedCh <- Counter{Bytes: bytes}:
	case <-p.closed:
	}
}

// CompleteItem is the status callback function for the archiver when a
// file/dir has been saved successfully.
func (p *Progress) CompleteItem(item string, previous, current *restic.Node, s archiver.ItemStats, d time.Duration) {
	p.summary.Lock()
	p.summary.ItemStats.Add(s)

	// for the last item "/", current is nil
	if current != nil {
		p.summary.ProcessedBytes += current.Size
	}

	p.summary.Unlock()

	if current == nil {
		// error occurred, tell the status display to remove the line
		select {
		case p.workerCh <- fileWorkerMessage{filename: item, done: true}:
		case <-p.closed:
		}
		return
	}

	switch current.Type {
	case "file":
		select {
		case p.processedCh <- Counter{Files: 1}:
		case <-p.closed:
		}
		select {
		case p.workerCh <- fileWorkerMessage{filename: item, done: true}:
		case <-p.closed:
		}
	case "dir":
		select {
		case p.processedCh <- Counter{Dirs: 1}:
		case <-p.closed:
		}
	}

	if current.Type == "dir" {
		if previous == nil {
			p.printer.CompleteItem("dir new", item, previous, current, s, d)
			p.summary.Lock()
			p.summary.Dirs.New++
			p.summary.Unlock()
			return
		}

		if previous.Equals(*current) {
			p.printer.CompleteItem("dir unchanged", item, previous, current, s, d)
			p.summary.Lock()
			p.summary.Dirs.Unchanged++
			p.summary.Unlock()
		} else {
			p.printer.CompleteItem("dir modified", item, previous, current, s, d)
			p.summary.Lock()
			p.summary.Dirs.Changed++
			p.summary.Unlock()
		}

	} else if current.Type == "file" {
		select {
		case p.workerCh <- fileWorkerMessage{done: true, filename: item}:
		case <-p.closed:
		}

		if previous == nil {
			p.printer.CompleteItem("file new", item, previous, current, s, d)
			p.summary.Lock()
			p.summary.Files.New++
			p.summary.Unlock()
			return
		}

		if previous.Equals(*current) {
			p.printer.CompleteItem("file unchanged", item, previous, current, s, d)
			p.summary.Lock()
			p.summary.Files.Unchanged++
			p.summary.Unlock()
		} else {
			p.printer.CompleteItem("file modified", item, previous, current, s, d)
			p.summary.Lock()
			p.summary.Files.Changed++
			p.summary.Unlock()
		}
	}
}

// ReportTotal sets the total stats up to now
func (p *Progress) ReportTotal(item string, s archiver.ScanStats) {
	select {
	case p.totalCh <- Counter{Files: uint64(s.Files), Dirs: uint64(s.Dirs), Bytes: s.Bytes}:
	case <-p.closed:
	}

	if item == "" {
		p.printer.ReportTotal(item, p.start, s)
		close(p.totalCh)
		return
	}
}

// Finish prints the finishing messages.
func (p *Progress) Finish(snapshotID restic.ID) {
	// wait for the status update goroutine to shut down
	<-p.closed
	p.printer.Finish(snapshotID, p.start, p.summary, p.dry)
}

// SetMinUpdatePause sets b.MinUpdatePause. It satisfies the
// ArchiveProgressReporter interface.
func (p *Progress) SetMinUpdatePause(d time.Duration) {
	p.MinUpdatePause = d
}

// SetDryRun marks the backup as a "dry run".
func (p *Progress) SetDryRun() {
	p.dry = true
}
<commit_msg>backup: Use buffered channels to collect backup status<commit_after>package backup

import (
	"context"
	"io"
	"sync"
	"time"

	"github.com/restic/restic/internal/archiver"
	"github.com/restic/restic/internal/restic"
	"github.com/restic/restic/internal/ui/signals"
)

type ProgressPrinter interface {
	Update(total, processed Counter, errors uint, currentFiles map[string]struct{}, start time.Time, secs uint64)
	Error(item string, err error) error
	ScannerError(item string, err error) error
	CompleteItem(messageType string, item string, previous, current *restic.Node, s archiver.ItemStats, d time.Duration)
	ReportTotal(item string, start time.Time, s archiver.ScanStats)
	Finish(snapshotID restic.ID, start time.Time, summary *Summary, dryRun bool)
	Reset()

	// ui.StdioWrapper
	Stdout() io.WriteCloser
	Stderr() io.WriteCloser

	E(msg string, args ...interface{})
	P(msg string, args ...interface{})
	V(msg string, args ...interface{})
	VV(msg string, args ...interface{})
}

type Counter struct {
	Files, Dirs, Bytes uint64
}

type fileWorkerMessage struct {
	filename string
	done     bool
}

type ProgressReporter interface {
	CompleteItem(item string, previous, current *restic.Node, s archiver.ItemStats, d time.Duration)
	StartFile(filename string)
	CompleteBlob(filename string, bytes uint64)
	ScannerError(item string, err error) error
	ReportTotal(item string, s archiver.ScanStats)
	SetMinUpdatePause(d time.Duration)
	Run(ctx context.Context) error
	Error(item string, err error) error
	Finish(snapshotID restic.ID)
}

type Summary struct {
	sync.Mutex
	Files, Dirs struct {
		New       uint
		Changed   uint
		Unchanged uint
	}
	ProcessedBytes uint64
	archiver.ItemStats
}

// Progress reports progress for the `backup` command.
type Progress struct {
	MinUpdatePause time.Duration

	start time.Time
	dry   bool

	totalBytes uint64

	totalCh     chan Counter
	processedCh chan Counter
	errCh       chan struct{}
	workerCh    chan fileWorkerMessage
	closed      chan struct{}

	summary *Summary
	printer ProgressPrinter
}

func NewProgress(printer ProgressPrinter) *Progress {
	return &Progress{
		// limit to 60fps by default
		MinUpdatePause: time.Second / 60,
		start:          time.Now(),

		// use buffered channels for the information used to update the status
		// the shutdown of the `Run()` method is somewhat racy, but won't affect
		// the final backup statistics
		totalCh:     make(chan Counter, 100),
		processedCh: make(chan Counter, 100),
		errCh:       make(chan struct{}),
		workerCh:    make(chan fileWorkerMessage, 100),
		closed:      make(chan struct{}),

		summary: &Summary{},

		printer: printer,
	}
}

// Run regularly updates the status lines. It should be called in a separate
// goroutine.
func (p *Progress) Run(ctx context.Context) error {
	var (
		lastUpdate       time.Time
		total, processed Counter
		errors           uint
		started          bool
		currentFiles     = make(map[string]struct{})
		secondsRemaining uint64
	)

	t := time.NewTicker(time.Second)
	signalsCh := signals.GetProgressChannel()
	defer t.Stop()
	defer close(p.closed)
	// Reset status when finished
	defer p.printer.Reset()

	for {
		forceUpdate := false
		select {
		case <-ctx.Done():
			return nil
		case t, ok := <-p.totalCh:
			if ok {
				total = t
				started = true
			} else {
				// scan has finished
				p.totalCh = nil
				p.totalBytes = total.Bytes
			}
		case s := <-p.processedCh:
			processed.Files += s.Files
			processed.Dirs += s.Dirs
			processed.Bytes += s.Bytes
			started = true
		case <-p.errCh:
			errors++
			started = true
		case m := <-p.workerCh:
			if m.done {
				delete(currentFiles, m.filename)
			} else {
				currentFiles[m.filename] = struct{}{}
			}
		case <-t.C:
			if !started {
				continue
			}

			if p.totalCh == nil {
				secs := float64(time.Since(p.start) / time.Second)
				todo := float64(total.Bytes - processed.Bytes)
				secondsRemaining = uint64(secs / float64(processed.Bytes) * todo)
			}
		case <-signalsCh:
			forceUpdate = true
		}

		// limit update frequency
		if !forceUpdate && (time.Since(lastUpdate) < p.MinUpdatePause || p.MinUpdatePause == 0) {
			continue
		}
		lastUpdate = time.Now()

		p.printer.Update(total, processed, errors, currentFiles, p.start, secondsRemaining)
	}
}

// ScannerError is the error callback function for the scanner, it prints the
// error in verbose mode and returns nil.
func (p *Progress) ScannerError(item string, err error) error {
	return p.printer.ScannerError(item, err)
}

// Error is the error callback function for the archiver, it prints the error and returns nil.
func (p *Progress) Error(item string, err error) error {
	cbErr := p.printer.Error(item, err)

	select {
	case p.errCh <- struct{}{}:
	case <-p.closed:
	}
	return cbErr
}

// StartFile is called when a file is being processed by a worker.
func (p *Progress) StartFile(filename string) {
	select {
	case p.workerCh <- fileWorkerMessage{filename: filename}:
	case <-p.closed:
	}
}

// CompleteBlob is called for all saved blobs for files.
func (p *Progress) CompleteBlob(filename string, bytes uint64) {
	select {
	case p.processedCh <- Counter{Bytes: bytes}:
	case <-p.closed:
	}
}

// CompleteItem is the status callback function for the archiver when a
// file/dir has been saved successfully.
func (p *Progress) CompleteItem(item string, previous, current *restic.Node, s archiver.ItemStats, d time.Duration) {
	p.summary.Lock()
	p.summary.ItemStats.Add(s)

	// for the last item "/", current is nil
	if current != nil {
		p.summary.ProcessedBytes += current.Size
	}

	p.summary.Unlock()

	if current == nil {
		// error occurred, tell the status display to remove the line
		select {
		case p.workerCh <- fileWorkerMessage{filename: item, done: true}:
		case <-p.closed:
		}
		return
	}

	switch current.Type {
	case "file":
		select {
		case p.processedCh <- Counter{Files: 1}:
		case <-p.closed:
		}
		select {
		case p.workerCh <- fileWorkerMessage{filename: item, done: true}:
		case <-p.closed:
		}
	case "dir":
		select {
		case p.processedCh <- Counter{Dirs: 1}:
		case <-p.closed:
		}
	}

	if current.Type == "dir" {
		if previous == nil {
			p.printer.CompleteItem("dir new", item, previous, current, s, d)
			p.summary.Lock()
			p.summary.Dirs.New++
			p.summary.Unlock()
			return
		}

		if previous.Equals(*current) {
			p.printer.CompleteItem("dir unchanged", item, previous, current, s, d)
			p.summary.Lock()
			p.summary.Dirs.Unchanged++
			p.summary.Unlock()
		} else {
			p.printer.CompleteItem("dir modified", item, previous, current, s, d)
			p.summary.Lock()
			p.summary.Dirs.Changed++
			p.summary.Unlock()
		}

	} else if current.Type == "file" {
		select {
		case p.workerCh <- fileWorkerMessage{done: true, filename: item}:
		case <-p.closed:
		}

		if previous == nil {
			p.printer.CompleteItem("file new", item, previous, current, s, d)
			p.summary.Lock()
			p.summary.Files.New++
			p.summary.Unlock()
			return
		}

		if previous.Equals(*current) {
			p.printer.CompleteItem("file unchanged", item, previous, current, s, d)
			p.summary.Lock()
			p.summary.Files.Unchanged++
			p.summary.Unlock()
		} else {
			p.printer.CompleteItem("file modified", item, previous, current, s, d)
			p.summary.Lock()
			p.summary.Files.Changed++
			p.summary.Unlock()
		}
	}
}

// ReportTotal sets the total stats up to now
func (p *Progress) ReportTotal(item string, s archiver.ScanStats) {
	select {
	case p.totalCh <- Counter{Files: uint64(s.Files), Dirs: uint64(s.Dirs), Bytes: s.Bytes}:
	case <-p.closed:
	}

	if item == "" {
		p.printer.ReportTotal(item, p.start, s)
		close(p.totalCh)
		return
	}
}

// Finish prints the finishing messages.
func (p *Progress) Finish(snapshotID restic.ID) {
	// wait for the status update goroutine to shut down
	<-p.closed
	p.printer.Finish(snapshotID, p.start, p.summary, p.dry)
}

// SetMinUpdatePause sets b.MinUpdatePause. It satisfies the
// ArchiveProgressReporter interface.
func (p *Progress) SetMinUpdatePause(d time.Duration) {
	p.MinUpdatePause = d
}

// SetDryRun marks the backup as a "dry run".
func (p *Progress) SetDryRun() {
	p.dry = true
}
<|endoftext|>


<commit_before>// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build go1.17
// +build go1.17

package worker

import (
	"context"
	"testing"
	"time"

	"github.com/go-git/go-git/v5"
	"github.com/go-git/go-git/v5/plumbing"
	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	"golang.org/x/vuln/internal/cveschema"
	"golang.org/x/vuln/internal/worker/log"
	"golang.org/x/vuln/internal/worker/store"
)

func TestRepoCVEFiles(t *testing.T) {
	repo, err := readTxtarRepo("testdata/basic.txtar", time.Now())
	if err != nil {
		t.Fatal(err)
	}
	commit := headCommit(t, repo)
	if err != nil {
		t.Fatal(err)
	}
	got, err := repoCVEFiles(repo, commit)
	if err != nil {
		t.Fatal(err)
	}

	want := []repoFile{
		{dirPath: "2020/9xxx", filename: "CVE-2020-9283.json", year: 2020, number: 9283},
		{dirPath: "2021/0xxx", filename: "CVE-2021-0001.json", year: 2021, number: 1},
		{dirPath: "2021/0xxx", filename: "CVE-2021-0010.json", year: 2021, number: 10},
		{dirPath: "2021/1xxx", filename: "CVE-2021-1384.json", year: 2021, number: 1384},
	}

	opt := cmpopts.IgnoreFields(repoFile{}, "treeHash", "blobHash")
	if diff := cmp.Diff(want, got, cmp.AllowUnexported(repoFile{}), opt); diff != "" {
		t.Errorf("mismatch (-want, +got):\n%s", diff)
	}
}

func TestDoUpdate(t *testing.T) {
	ctx := log.WithLineLogger(context.Background())
	repo, err := readTxtarRepo("testdata/basic.txtar", time.Now())
	if err != nil {
		t.Fatal(err)
	}
	h, err := headHash(repo)
	if err != nil {
		t.Fatal(err)
	}

	purl := getPkgsiteURL(t)
	needsIssue := func(cve *cveschema.CVE) (*triageResult, error) {
		return TriageCVE(ctx, cve, purl)
	}

	ref, err := repo.Reference(plumbing.HEAD, true)
	if err != nil {
		t.Fatal(err)
	}

	commitHash := ref.Hash().String()
	knownVulns := []string{"CVE-2020-9283"}

	paths := []string{
		"2021/0xxx/CVE-2021-0001.json",
		"2021/0xxx/CVE-2021-0010.json",
		"2021/1xxx/CVE-2021-1384.json",
		"2020/9xxx/CVE-2020-9283.json",
	}

	var (
		cves       []*cveschema.CVE
		blobHashes []string
	)
	for _, p := range paths {
		cve, bh := readCVE(t, repo, p)
		cves = append(cves, cve)
		blobHashes = append(blobHashes, bh)
	}
	// CVERecords after the above CVEs are added to an empty DB.
	var rs []*store.CVERecord
	for i := 0; i < len(cves); i++ {
		r := &store.CVERecord{
			ID:         cves[i].ID,
			CVEState:   cves[i].State,
			Path:       paths[i],
			BlobHash:   blobHashes[i],
			CommitHash: commitHash,
		}
		rs = append(rs, r)
	}
	rs[0].TriageState = store.TriageStateNeedsIssue // a public CVE, has a golang.org path
	rs[0].Module = "golang.org/x/mod"
	rs[0].CVE = cves[0]
	rs[1].TriageState = store.TriageStateNoActionNeeded // state is reserved
	rs[2].TriageState = store.TriageStateNoActionNeeded // state is rejected
	rs[3].TriageState = store.TriageStateNoActionNeeded
	rs[3].TriageStateReason = "already in vuln DB"

	// withTriageState returns a copy of r with the TriageState field changed to ts.
	withTriageState := func(r *store.CVERecord, ts store.TriageState) *store.CVERecord {
		c := *r
		c.BlobHash += "x" // if we don't use a different blob hash, no update will happen
		c.CommitHash = "?"
		c.TriageState = ts
		return &c
	}

	for _, test := range []struct {
		name string
		cur  []*store.CVERecord // current state of DB
		want []*store.CVERecord // expected state after update
	}{
		{
			name: "empty",
			cur:  nil,
			want: rs,
		},
		{
			name: "no change",
			cur:  rs,
			want: rs,
		},
		{
			name: "pre-issue changes",
			cur: []*store.CVERecord{
				// NoActionNeeded -> NeedsIssue
				withTriageState(rs[0], store.TriageStateNoActionNeeded),
				// NeedsIssue -> NoActionNeeded
				func() *store.CVERecord {
					r := withTriageState(rs[1], store.TriageStateNeedsIssue)
					r.Module = "something"
					r.CVE = cves[1]
					return r
				}(),
				// NoActionNeeded, triage state stays the same but other fields change.
				withTriageState(rs[2], store.TriageStateNoActionNeeded),
			},
			want: []*store.CVERecord{
				rs[0],
				func() *store.CVERecord {
					c := *rs[1]
					c.Module = ""
					c.CVE = nil
					return &c
				}(),
				rs[2],
				rs[3],
			},
		},
		{
			name: "post-issue changes",
			cur: []*store.CVERecord{
				// IssueCreated -> Updated
				withTriageState(rs[0], store.TriageStateIssueCreated),
				withTriageState(rs[1], store.TriageStateUpdatedSinceIssueCreation),
			},
			want: []*store.CVERecord{
				func() *store.CVERecord {
					c := *rs[0]
					c.TriageState = store.TriageStateUpdatedSinceIssueCreation
					c.TriageStateReason = `CVE changed; affected module = "golang.org/x/mod"`
					return &c
				}(),
				func() *store.CVERecord {
					c := *rs[1]
					c.TriageState = store.TriageStateUpdatedSinceIssueCreation
					c.TriageStateReason = `CVE changed; affected module = ""`
					return &c
				}(),
				rs[2],
				rs[3],
			},
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			mstore := store.NewMemStore()
			createCVERecords(t, mstore, test.cur)
			if _, err := newUpdater(repo, h, mstore, knownVulns, needsIssue).update(ctx); err != nil {
				t.Fatal(err)
			}
			got := mstore.CVERecords()
			want := map[string]*store.CVERecord{}
			for _, cr := range test.want {
				want[cr.ID] = cr
			}
			if diff := cmp.Diff(want, got); diff != "" {
				t.Errorf("mismatch (-want, +got):\n%s", diff)
			}
		})
	}
}

func TestGroupFilesByDirectory(t *testing.T) {
	for _, test := range []struct {
		in   []repoFile
		want [][]repoFile
	}{
		{in: nil, want: nil},
		{
			in:   []repoFile{{dirPath: "a"}},
			want: [][]repoFile{{{dirPath: "a"}}},
		},
		{
			in: []repoFile{
				{dirPath: "a", filename: "f1"},
				{dirPath: "a", filename: "f2"},
			},
			want: [][]repoFile{{
				{dirPath: "a", filename: "f1"},
				{dirPath: "a", filename: "f2"},
			}},
		},
		{
			in: []repoFile{
				{dirPath: "a", filename: "f1"},
				{dirPath: "a", filename: "f2"},
				{dirPath: "b", filename: "f1"},
				{dirPath: "c", filename: "f1"},
				{dirPath: "c", filename: "f2"},
			},
			want: [][]repoFile{
				{
					{dirPath: "a", filename: "f1"},
					{dirPath: "a", filename: "f2"},
				},
				{
					{dirPath: "b", filename: "f1"},
				},
				{
					{dirPath: "c", filename: "f1"},
					{dirPath: "c", filename: "f2"},
				},
			},
		},
	} {
		got, err := groupFilesByDirectory(test.in)
		if err != nil {
			t.Fatalf("%v: %v", test.in, err)
		}
		if diff := cmp.Diff(got, test.want, cmp.AllowUnexported(repoFile{})); diff != "" {
			t.Errorf("%v: (-want, +got)\n%s", test.in, diff)
		}
	}

	_, err := groupFilesByDirectory([]repoFile{{dirPath: "a"}, {dirPath: "b"}, {dirPath: "a"}})
	if err == nil {
		t.Error("got nil, want error")
	}
}

func readCVE(t *testing.T, repo *git.Repository, path string) (*cveschema.CVE, string) {
	c := headCommit(t, repo)
	cve, blobHash, err := ReadCVEAtPath(c, path)
	if err != nil {
		t.Fatal(err)
	}
	return cve, blobHash
}

func createCVERecords(t *testing.T, s store.Store, crs []*store.CVERecord) {
	err := s.RunTransaction(context.Background(), func(_ context.Context, tx store.Transaction) error {
		for _, cr := range crs {
			if err := tx.CreateCVERecord(cr); err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		t.Fatal(err)
	}
}
<commit_msg>internal/worker: improve update test<commit_after>// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build go1.17
// +build go1.17

package worker

import (
	"context"
	"testing"
	"time"

	"github.com/go-git/go-git/v5"
	"github.com/go-git/go-git/v5/plumbing"
	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	"golang.org/x/vuln/internal/cveschema"
	"golang.org/x/vuln/internal/worker/log"
	"golang.org/x/vuln/internal/worker/store"
)

func TestRepoCVEFiles(t *testing.T) {
	repo, err := readTxtarRepo("testdata/basic.txtar", time.Now())
	if err != nil {
		t.Fatal(err)
	}
	commit := headCommit(t, repo)
	if err != nil {
		t.Fatal(err)
	}
	got, err := repoCVEFiles(repo, commit)
	if err != nil {
		t.Fatal(err)
	}

	want := []repoFile{
		{dirPath: "2020/9xxx", filename: "CVE-2020-9283.json", year: 2020, number: 9283},
		{dirPath: "2021/0xxx", filename: "CVE-2021-0001.json", year: 2021, number: 1},
		{dirPath: "2021/0xxx", filename: "CVE-2021-0010.json", year: 2021, number: 10},
		{dirPath: "2021/1xxx", filename: "CVE-2021-1384.json", year: 2021, number: 1384},
	}

	opt := cmpopts.IgnoreFields(repoFile{}, "treeHash", "blobHash")
	if diff := cmp.Diff(want, got, cmp.AllowUnexported(repoFile{}), opt); diff != "" {
		t.Errorf("mismatch (-want, +got):\n%s", diff)
	}
}

const clearString = "**CLEAR**"

var clearCVE = &cveschema.CVE{}

func modify(r, m *store.CVERecord) *store.CVERecord {
	modString := func(p *string, s string) {
		if s == clearString {
			*p = ""
		} else if s != "" {
			*p = s
		}
	}

	c := *r
	modString(&c.BlobHash, m.BlobHash)
	modString(&c.CommitHash, m.CommitHash)
	modString(&c.CVEState, m.CVEState)
	if m.TriageState != "" {
		if m.TriageState == clearString {
			c.TriageState = ""
		} else {
			c.TriageState = m.TriageState
		}
	}
	modString(&c.TriageStateReason, m.TriageStateReason)
	modString(&c.Module, m.Module)
	if m.CVE == clearCVE {
		c.CVE = nil
	} else if m.CVE != nil {
		c.CVE = m.CVE
	}
	modString(&c.IssueReference, m.IssueReference)
	if !m.IssueCreatedAt.IsZero() {
		panic("unsupported modification")
	}
	return &c
}

func TestDoUpdate(t *testing.T) {
	ctx := log.WithLineLogger(context.Background())
	repo, err := readTxtarRepo("testdata/basic.txtar", time.Now())
	if err != nil {
		t.Fatal(err)
	}
	h, err := headHash(repo)
	if err != nil {
		t.Fatal(err)
	}

	purl := getPkgsiteURL(t)
	needsIssue := func(cve *cveschema.CVE) (*triageResult, error) {
		return TriageCVE(ctx, cve, purl)
	}

	ref, err := repo.Reference(plumbing.HEAD, true)
	if err != nil {
		t.Fatal(err)
	}

	commitHash := ref.Hash().String()
	knownVulns := []string{"CVE-2020-9283"}

	paths := []string{
		"2021/0xxx/CVE-2021-0001.json",
		"2021/0xxx/CVE-2021-0010.json",
		"2021/1xxx/CVE-2021-1384.json",
		"2020/9xxx/CVE-2020-9283.json",
	}

	var (
		cves       []*cveschema.CVE
		blobHashes []string
	)
	for _, p := range paths {
		cve, bh := readCVE(t, repo, p)
		cves = append(cves, cve)
		blobHashes = append(blobHashes, bh)
	}
	// CVERecords after the above CVEs are added to an empty DB.
	var rs []*store.CVERecord
	for i := 0; i < len(cves); i++ {
		r := &store.CVERecord{
			ID:         cves[i].ID,
			CVEState:   cves[i].State,
			Path:       paths[i],
			BlobHash:   blobHashes[i],
			CommitHash: commitHash,
		}
		rs = append(rs, r)
	}
	rs[0].TriageState = store.TriageStateNeedsIssue // a public CVE, has a golang.org path
	rs[0].Module = "golang.org/x/mod"
	rs[0].CVE = cves[0]
	rs[1].TriageState = store.TriageStateNoActionNeeded // state is reserved
	rs[2].TriageState = store.TriageStateNoActionNeeded // state is rejected
	rs[3].TriageState = store.TriageStateNoActionNeeded
	rs[3].TriageStateReason = "already in vuln DB"

	for _, test := range []struct {
		name string
		cur  []*store.CVERecord // current state of DB
		want []*store.CVERecord // expected state after update
	}{
		{
			name: "empty",
			cur:  nil,
			want: rs,
		},
		{
			name: "no change",
			cur:  rs,
			want: rs,
		},
		{
			name: "pre-issue changes",
			cur: []*store.CVERecord{
				// NoActionNeeded -> NeedsIssue
				modify(rs[0], &store.CVERecord{
					BlobHash:    "x", // if we don't use a different blob hash, no update will happen
					TriageState: store.TriageStateNoActionNeeded,
				}),
				// NeedsIssue -> NoActionNeeded
				modify(rs[1], &store.CVERecord{
					BlobHash:    "x",
					TriageState: store.TriageStateNeedsIssue,
					Module:      "something",
					CVE:         cves[1],
				}),
				// NoActionNeeded, triage state stays the same but other fields change.
				modify(rs[2], &store.CVERecord{
					TriageState: store.TriageStateNoActionNeeded,
				}),
			},
			want: []*store.CVERecord{
				rs[0],
				modify(rs[1], &store.CVERecord{
					Module: clearString,
					CVE:    clearCVE,
				}),
				rs[2],
				rs[3],
			},
		},
		{
			name: "post-issue changes",
			cur: []*store.CVERecord{
				// IssueCreated -> Updated
				modify(rs[0], &store.CVERecord{
					BlobHash:    "x",
					TriageState: store.TriageStateIssueCreated,
				}),
				modify(rs[1], &store.CVERecord{
					BlobHash:    "x",
					TriageState: store.TriageStateUpdatedSinceIssueCreation,
				}),
			},
			want: []*store.CVERecord{
				modify(rs[0], &store.CVERecord{
					TriageState:       store.TriageStateUpdatedSinceIssueCreation,
					TriageStateReason: `CVE changed; affected module = "golang.org/x/mod"`,
				}),
				modify(rs[1], &store.CVERecord{
					TriageState:       store.TriageStateUpdatedSinceIssueCreation,
					TriageStateReason: `CVE changed; affected module = ""`,
				}),
				rs[2],
				rs[3],
			},
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			mstore := store.NewMemStore()
			createCVERecords(t, mstore, test.cur)
			if _, err := newUpdater(repo, h, mstore, knownVulns, needsIssue).update(ctx); err != nil {
				t.Fatal(err)
			}
			got := mstore.CVERecords()
			want := map[string]*store.CVERecord{}
			for _, cr := range test.want {
				want[cr.ID] = cr
			}
			if diff := cmp.Diff(want, got); diff != "" {
				t.Errorf("mismatch (-want, +got):\n%s", diff)
			}
		})
	}
}

func TestGroupFilesByDirectory(t *testing.T) {
	for _, test := range []struct {
		in   []repoFile
		want [][]repoFile
	}{
		{in: nil, want: nil},
		{
			in:   []repoFile{{dirPath: "a"}},
			want: [][]repoFile{{{dirPath: "a"}}},
		},
		{
			in: []repoFile{
				{dirPath: "a", filename: "f1"},
				{dirPath: "a", filename: "f2"},
			},
			want: [][]repoFile{{
				{dirPath: "a", filename: "f1"},
				{dirPath: "a", filename: "f2"},
			}},
		},
		{
			in: []repoFile{
				{dirPath: "a", filename: "f1"},
				{dirPath: "a", filename: "f2"},
				{dirPath: "b", filename: "f1"},
				{dirPath: "c", filename: "f1"},
				{dirPath: "c", filename: "f2"},
			},
			want: [][]repoFile{
				{
					{dirPath: "a", filename: "f1"},
					{dirPath: "a", filename: "f2"},
				},
				{
					{dirPath: "b", filename: "f1"},
				},
				{
					{dirPath: "c", filename: "f1"},
					{dirPath: "c", filename: "f2"},
				},
			},
		},
	} {
		got, err := groupFilesByDirectory(test.in)
		if err != nil {
			t.Fatalf("%v: %v", test.in, err)
		}
		if diff := cmp.Diff(got, test.want, cmp.AllowUnexported(repoFile{})); diff != "" {
			t.Errorf("%v: (-want, +got)\n%s", test.in, diff)
		}
	}

	_, err := groupFilesByDirectory([]repoFile{{dirPath: "a"}, {dirPath: "b"}, {dirPath: "a"}})
	if err == nil {
		t.Error("got nil, want error")
	}
}

func readCVE(t *testing.T, repo *git.Repository, path string) (*cveschema.CVE, string) {
	c := headCommit(t, repo)
	cve, blobHash, err := ReadCVEAtPath(c, path)
	if err != nil {
		t.Fatal(err)
	}
	return cve, blobHash
}

func createCVERecords(t *testing.T, s store.Store, crs []*store.CVERecord) {
	err := s.RunTransaction(context.Background(), func(_ context.Context, tx store.Transaction) error {
		for _, cr := range crs {
			if err := tx.CreateCVERecord(cr); err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		t.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package internal

import (
	"encoding/xml"
	"io"

	"github.com/szyhf/go-convert"
)

func readSharedStringsXML(rc io.ReadCloser) []string {
	decoder := xml.NewDecoder(rc)

	tStart, rStart := false, false
	var slc []string
	siIndex := 0
	for t, err := decoder.Token(); err == nil; t, err = decoder.Token() {
		switch token := t.(type) {
		case xml.StartElement:
			switch token.Name.Local {
			case "si":
				// don't enter default ...
			case "t":
				tStart = true
			case "r":
				rStart = true
			case "sst":
				for _, attr := range token.Attr {
					if attr.Name.Local == "uniqueCount" {
						slc = make([]string, convert.MustInt(attr.Value))
					}
				}
			default:
				decoder.Skip()
			}
		case xml.EndElement:
			switch token.Name.Local {
			case "si":
				siIndex++
			case "t":
				tStart = false
			case "r":
				rStart = false
			}
		case xml.CharData:
			if tStart {
				if rStart {
					str := slc[siIndex]
					slc[siIndex] = str + string(token)
				} else {
					slc[siIndex] = string(token)
				}
			}
		}
	}
	return slc
}
<commit_msg>嗯。。sharedString…<commit_after>package internal

import (
	"encoding/xml"
	"io"

	"github.com/szyhf/go-convert"
)

func readSharedStringsXML(rc io.ReadCloser) []string {
	decoder := xml.NewDecoder(rc)

	tStart, rStart := false, false
	var slc []string
	siIndex := 0
	for t, err := decoder.Token(); err == nil; t, err = decoder.Token() {
		switch token := t.(type) {
		case xml.StartElement:
			switch token.Name.Local {
			case "si":
				// don't enter default ...
			case "t":
				tStart = true
			case "r":
				rStart = true
			case "sst":
				count := 0
				unqCount := 0
				for _, attr := range token.Attr {
					switch attr.Name.Local {
					case "count":
						count = convert.MustInt(attr.Value)
					case "uniqueCount":
						unqCount = convert.MustInt(attr.Value)
					}
				}
				if unqCount != 0 {
					slc = make([]string, unqCount)
				} else {
					slc = make([]string, count)
				}
			default:
				decoder.Skip()
			}
		case xml.EndElement:
			switch token.Name.Local {
			case "si":
				siIndex++
			case "t":
				tStart = false
			case "r":
				rStart = false
			}
		case xml.CharData:
			if tStart {
				if rStart {
					str := slc[siIndex]
					slc[siIndex] = str + string(token)
				} else {
					slc[siIndex] = string(token)
				}
			}
		}
	}
	return slc
}
<|endoftext|>


<commit_before>// Package flagconfig provides a flag to specifiy a config file which
// will in turn be used to read unspecified flag values.
package flagconfig

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
)

var configFile = flag.String(
	"c", defaultConfig(), "Config file to read flags from.")

func Usage() {
	Parse()
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	flag.VisitAll(func(f *flag.Flag) {
		fmt.Fprintf(os.Stderr, "  -%s=%s: %s\n", f.Name, f.Value.String(), f.Usage)
	})
}

func defaultConfig() string {
	home := os.Getenv("HOME")
	basename := filepath.Base(os.Args[0])
	path := filepath.Join(home, ".config", basename, "config")
	_, err := os.Open(path)
	if err == nil {
		return path
	}
	path = filepath.Join("/", "etc", "conf.d", basename)
	_, err = os.Open(path)
	if err == nil {
		return path
	}
	return ""
}

func contains(list []*flag.Flag, f *flag.Flag) bool {
	for _, i := range list {
		if i == f {
			return true
		}
	}
	return false
}

func readConfig(filename string) map[string]string {
	bytes, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Fatalf("Failed to read config file %s: %s", *configFile, err)
	}
	lines := strings.Split(string(bytes), "\n")
	result := make(map[string]string, len(lines))
	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if trimmed == "" || trimmed[0] == '#' {
			continue
		}
		parts := strings.Split(line, "=")
		if len(parts) != 2 {
			log.Fatalf("Invalid config line: %s", line)
		}
		result[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1])
	}
	return result
}

func Parse() {
	ParseSet(flag.CommandLine)
}

func ParseSet(set *flag.FlagSet) {
	ParseFile(set, *configFile)
}

func ParseFile(set *flag.FlagSet, filename string) {
	if filename == "" {
		return
	}
	config := readConfig(filename)
	explicit := make([]*flag.Flag, 0)
	all := make([]*flag.Flag, 0)
	set.Visit(func(f *flag.Flag) {
		explicit = append(explicit, f)
	})
	set.VisitAll(func(f *flag.Flag) {
		all = append(all, f)
		if !contains(explicit, f) {
			val := config[f.Name]
			if val != "" {
				err := f.Value.Set(val)
				if err != nil {
					log.Fatalf("Failed to set flag %s with value %s", f.Name, val)
				}
			}
		}
	})
Outer:
	for name, val := range config {
		for _, f := range all {
			if f.Name == name {
				continue Outer
			}
		}
		log.Fatalf("Unknown flag %s=%s in config file.", name, val)
	}
}
<commit_msg>Added documentation and a couple small refactorings (#3)<commit_after>// Package flagconfig provides a flag to specifiy a config file which
// will in turn be used to read unspecified flag values.
//
// Default configuration file location is $HOME/.config/{executable_name}/config
// or /etc/conf.d/{executable_name}
//
// Lines in the configuration file should be in flag-name=value format.
// Comments are allowed in the configuration file.
package flagconfig

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
)

var configFile = flag.String(
	"c", defaultConfig(), "Config file to read flags from.")

// Usage prints the usage information for the application including all flags
// and their values after parsing the configuration file
func Usage() {
	Parse()
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	flag.VisitAll(func(f *flag.Flag) {
		fmt.Fprintf(os.Stderr, "  -%s=%s: %s\n", f.Name, f.Value.String(), f.Usage)
	})
}

func defaultConfig() string {
	home := os.Getenv("HOME")
	basename := filepath.Base(os.Args[0])
	path := filepath.Join(home, ".config", basename, "config")
	_, err := os.Open(path)
	if err == nil {
		return path
	}
	path = filepath.Join("/", "etc", "conf.d", basename)
	_, err = os.Open(path)
	if err == nil {
		return path
	}
	return ""
}

func contains(list []*flag.Flag, f *flag.Flag) bool {
	for _, i := range list {
		if i == f {
			return true
		}
	}
	return false
}

func readConfig(filename string) map[string]string {
	bytes, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Fatalf("Failed to read config file %s: %s", *configFile, err)
	}
	lines := strings.Split(string(bytes), "\n")
	result := make(map[string]string, len(lines))
	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if trimmed == "" || trimmed[0] == '#' {
			continue
		}
		parts := strings.Split(line, "=")
		if len(parts) != 2 {
			log.Fatalf("Invalid config line: %s", line)
		}
		result[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1])
	}
	return result
}

// Parse parses the default configuration file and populates the global flags
// based on the contents of the file.
func Parse() {
	ParseSet(flag.CommandLine)
}

// ParseSet parses the default configuraiton file and populates the flags in
// the flag.FlagSet based on the contents of the file.
func ParseSet(flags *flag.FlagSet) {
	ParseFile(flags, *configFile)
}

// ParseFile parses the specified configuration file and populates the flags
// in the flag.FlagSet based on the contents of the file.
func ParseFile(flags *flag.FlagSet, filename string) {
	if filename == "" {
		return
	}

	var (
		explicit []*flag.Flag
		all      []*flag.Flag
	)

	config := readConfig(filename)

	flags.Visit(func(f *flag.Flag) {
		explicit = append(explicit, f)
	})

	flags.VisitAll(func(f *flag.Flag) {
		all = append(all, f)
		if !contains(explicit, f) {
			val := config[f.Name]
			if val != "" {
				err := f.Value.Set(val)
				if err != nil {
					log.Fatalf("Failed to set flag %s with value %s", f.Name, val)
				}
			}
		}
	})
Outer:
	for name, val := range config {
		for _, f := range all {
			if f.Name == name {
				continue Outer
			}
		}
		log.Fatalf("Unknown flag %s=%s in config file.", name, val)
	}
}
<|endoftext|>


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

package main

import (
	"fmt"
	"os"
	"sort"
	"testing"
	"time"

	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"github.com/urfave/cli"
)

func TestFlags(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type FlagsTest struct {
}

func init() { RegisterTestSuite(&FlagsTest{}) }

func parseArgs(args []string) (flags *flagStorage) {
	// Create a CLI app, and abuse it to snoop on the flags.
	app := newApp()
	app.Action = func(appCtx *cli.Context) {
		flags = populateFlags(appCtx)
	}

	// Simulate argv.
	fullArgs := append([]string{"some_app"}, args...)

	err := app.Run(fullArgs)
	AssertEq(nil, err)

	return
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *FlagsTest) Defaults() {
	f := parseArgs([]string{})

	// File system
	ExpectNe(nil, f.MountOptions)
	ExpectEq(0, len(f.MountOptions), "Options: %v", f.MountOptions)

	ExpectEq(os.FileMode(0755), f.DirMode)
	ExpectEq(os.FileMode(0644), f.FileMode)
	ExpectEq(-1, f.Uid)
	ExpectEq(-1, f.Gid)
	ExpectFalse(f.ImplicitDirs)

	// GCS
	ExpectEq("", f.KeyFile)
	ExpectEq(-1, f.EgressBandwidthLimitBytesPerSecond)
	ExpectEq(-1, f.OpRateLimitHz)

	// Tuning
	ExpectEq(4096, f.StatCacheCapacity)
	ExpectEq(time.Minute, f.StatCacheTTL)
	ExpectEq(time.Minute, f.TypeCacheTTL)
	ExpectEq("", f.TempDir)

	// Debugging
	ExpectFalse(f.DebugFuse)
	ExpectFalse(f.DebugGCS)
	ExpectFalse(f.DebugHTTP)
	ExpectFalse(f.DebugInvariants)
}

func (t *FlagsTest) Bools() {
	names := []string{
		"implicit-dirs",
		"debug_fuse",
		"debug_gcs",
		"debug_http",
		"debug_invariants",
	}

	var args []string
	var f *flagStorage

	// --foo form
	args = nil
	for _, n := range names {
		args = append(args, fmt.Sprintf("-%s", n))
	}

	f = parseArgs(args)
	ExpectTrue(f.ImplicitDirs)
	ExpectTrue(f.DebugFuse)
	ExpectTrue(f.DebugGCS)
	ExpectTrue(f.DebugHTTP)
	ExpectTrue(f.DebugInvariants)

	// --foo=false form
	args = nil
	for _, n := range names {
		args = append(args, fmt.Sprintf("-%s=false", n))
	}

	f = parseArgs(args)
	ExpectFalse(f.ImplicitDirs)
	ExpectFalse(f.DebugFuse)
	ExpectFalse(f.DebugGCS)
	ExpectFalse(f.DebugHTTP)
	ExpectFalse(f.DebugInvariants)

	// --foo=true form
	args = nil
	for _, n := range names {
		args = append(args, fmt.Sprintf("-%s=true", n))
	}

	f = parseArgs(args)
	ExpectTrue(f.ImplicitDirs)
	ExpectTrue(f.DebugFuse)
	ExpectTrue(f.DebugGCS)
	ExpectTrue(f.DebugHTTP)
	ExpectTrue(f.DebugInvariants)
}

func (t *FlagsTest) DecimalNumbers() {
	args := []string{
		"--uid=17",
		"--gid=19",
		"--limit-bytes-per-sec=123.4",
		"--limit-ops-per-sec=56.78",
		"--stat-cache-capacity=8192",
	}

	f := parseArgs(args)
	ExpectEq(17, f.Uid)
	ExpectEq(19, f.Gid)
	ExpectEq(123.4, f.EgressBandwidthLimitBytesPerSecond)
	ExpectEq(56.78, f.OpRateLimitHz)
	ExpectEq(8192, f.StatCacheCapacity)
}

func (t *FlagsTest) OctalNumbers() {
	args := []string{
		"--dir-mode=711",
		"--file-mode", "611",
	}

	f := parseArgs(args)
	ExpectEq(os.FileMode(0711), f.DirMode)
	ExpectEq(os.FileMode(0611), f.FileMode)
}

func (t *FlagsTest) Strings() {
	args := []string{
		"--key-file", "-asdf",
		"--temp-dir=foobar",
		"--only-dir=baz",
	}

	f := parseArgs(args)
	ExpectEq("-asdf", f.KeyFile)
	ExpectEq("foobar", f.TempDir)
	ExpectEq("baz", f.OnlyDir)
}

func (t *FlagsTest) Durations() {
	args := []string{
		"--stat-cache-ttl", "1m17s",
		"--type-cache-ttl", "19ns",
	}

	f := parseArgs(args)
	ExpectEq(77*time.Second, f.StatCacheTTL)
	ExpectEq(19*time.Nanosecond, f.TypeCacheTTL)
}

func (t *FlagsTest) Maps() {
	args := []string{
		"-o", "rw,nodev",
		"-o", "user=jacobsa,noauto",
	}

	f := parseArgs(args)

	var keys sort.StringSlice
	for k := range f.MountOptions {
		keys = append(keys, k)
	}

	sort.Sort(keys)
	AssertThat(keys, ElementsAre("noauto", "nodev", "rw", "user"))

	ExpectEq("", f.MountOptions["noauto"])
	ExpectEq("", f.MountOptions["nodev"])
	ExpectEq("", f.MountOptions["rw"])
	ExpectEq("jacobsa", f.MountOptions["user"])
}
<commit_msg>Adding the unit test for ResolvePath method<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"testing"
	"time"

	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"github.com/urfave/cli"
)

func TestFlags(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type FlagsTest struct {
}

func init() { RegisterTestSuite(&FlagsTest{}) }

func parseArgs(args []string) (flags *flagStorage) {
	// Create a CLI app, and abuse it to snoop on the flags.
	app := newApp()
	app.Action = func(appCtx *cli.Context) {
		resolvePathForTheRequiredFlagInContext(appCtx)
		flags = populateFlags(appCtx)
	}

	// Simulate argv.
	fullArgs := append([]string{"some_app"}, args...)

	err := app.Run(fullArgs)
	AssertEq(nil, err)

	return
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *FlagsTest) Defaults() {
	f := parseArgs([]string{})

	// File system
	ExpectNe(nil, f.MountOptions)
	ExpectEq(0, len(f.MountOptions), "Options: %v", f.MountOptions)

	ExpectEq(os.FileMode(0755), f.DirMode)
	ExpectEq(os.FileMode(0644), f.FileMode)
	ExpectEq(-1, f.Uid)
	ExpectEq(-1, f.Gid)
	ExpectFalse(f.ImplicitDirs)

	// GCS
	ExpectEq("", f.KeyFile)
	ExpectEq(-1, f.EgressBandwidthLimitBytesPerSecond)
	ExpectEq(-1, f.OpRateLimitHz)

	// Tuning
	ExpectEq(4096, f.StatCacheCapacity)
	ExpectEq(time.Minute, f.StatCacheTTL)
	ExpectEq(time.Minute, f.TypeCacheTTL)
	ExpectEq("", f.TempDir)

	// Debugging
	ExpectFalse(f.DebugFuse)
	ExpectFalse(f.DebugGCS)
	ExpectFalse(f.DebugHTTP)
	ExpectFalse(f.DebugInvariants)
}

func (t *FlagsTest) Bools() {
	names := []string{
		"implicit-dirs",
		"debug_fuse",
		"debug_gcs",
		"debug_http",
		"debug_invariants",
	}

	var args []string
	var f *flagStorage

	// --foo form
	args = nil
	for _, n := range names {
		args = append(args, fmt.Sprintf("-%s", n))
	}

	f = parseArgs(args)
	ExpectTrue(f.ImplicitDirs)
	ExpectTrue(f.DebugFuse)
	ExpectTrue(f.DebugGCS)
	ExpectTrue(f.DebugHTTP)
	ExpectTrue(f.DebugInvariants)

	// --foo=false form
	args = nil
	for _, n := range names {
		args = append(args, fmt.Sprintf("-%s=false", n))
	}

	f = parseArgs(args)
	ExpectFalse(f.ImplicitDirs)
	ExpectFalse(f.DebugFuse)
	ExpectFalse(f.DebugGCS)
	ExpectFalse(f.DebugHTTP)
	ExpectFalse(f.DebugInvariants)

	// --foo=true form
	args = nil
	for _, n := range names {
		args = append(args, fmt.Sprintf("-%s=true", n))
	}

	f = parseArgs(args)
	ExpectTrue(f.ImplicitDirs)
	ExpectTrue(f.DebugFuse)
	ExpectTrue(f.DebugGCS)
	ExpectTrue(f.DebugHTTP)
	ExpectTrue(f.DebugInvariants)
}

func (t *FlagsTest) DecimalNumbers() {
	args := []string{
		"--uid=17",
		"--gid=19",
		"--limit-bytes-per-sec=123.4",
		"--limit-ops-per-sec=56.78",
		"--stat-cache-capacity=8192",
	}

	f := parseArgs(args)
	ExpectEq(17, f.Uid)
	ExpectEq(19, f.Gid)
	ExpectEq(123.4, f.EgressBandwidthLimitBytesPerSecond)
	ExpectEq(56.78, f.OpRateLimitHz)
	ExpectEq(8192, f.StatCacheCapacity)
}

func (t *FlagsTest) OctalNumbers() {
	args := []string{
		"--dir-mode=711",
		"--file-mode", "611",
	}

	f := parseArgs(args)
	ExpectEq(os.FileMode(0711), f.DirMode)
	ExpectEq(os.FileMode(0611), f.FileMode)
}

func (t *FlagsTest) Strings() {
	args := []string{
		"--key-file", "-asdf",
		"--temp-dir=foobar",
		"--only-dir=baz",
	}

	currentWorkingDir, err := os.Getwd()
	AssertEq(nil, err)

	f := parseArgs(args)
	ExpectEq(filepath.Join(currentWorkingDir, "-asdf"), f.KeyFile)
	ExpectEq("foobar", f.TempDir)
	ExpectEq("baz", f.OnlyDir)
}

func (t *FlagsTest) Durations() {
	args := []string{
		"--stat-cache-ttl", "1m17s",
		"--type-cache-ttl", "19ns",
	}

	f := parseArgs(args)
	ExpectEq(77*time.Second, f.StatCacheTTL)
	ExpectEq(19*time.Nanosecond, f.TypeCacheTTL)
}

func (t *FlagsTest) Maps() {
	args := []string{
		"-o", "rw,nodev",
		"-o", "user=jacobsa,noauto",
	}

	f := parseArgs(args)

	var keys sort.StringSlice
	for k := range f.MountOptions {
		keys = append(keys, k)
	}

	sort.Sort(keys)
	AssertThat(keys, ElementsAre("noauto", "nodev", "rw", "user"))

	ExpectEq("", f.MountOptions["noauto"])
	ExpectEq("", f.MountOptions["nodev"])
	ExpectEq("", f.MountOptions["rw"])
	ExpectEq("jacobsa", f.MountOptions["user"])
}

func (t *FlagsTest) ResolvePathInParentProcess() {
	args := []string{
		"--log-file=~/test.txt",
		"--key-file=test.json",
	}

	f := parseArgs(args)
	currentWorkingDir, err := os.Getwd()
	AssertEq(nil, err)

	homeDir, err := os.UserHomeDir()
	AssertEq(nil, err)

	ExpectEq(filepath.Join(homeDir, "test.txt"), f.LogFile)
	ExpectEq(filepath.Join(currentWorkingDir, "test.json"), f.KeyFile)
}

func (t *FlagsTest) ResolvePathInChildProcess() {
	args := []string{
		"--log-file=test.txt",
		"--key-file=test.json",
	}

	const gcsFuseParentProcessDir = "/var/generic/google"

	// By setting this environment variable, resolve will work for child process.
	os.Setenv(GCSFUSE_PARENT_PROCESS_DIR, gcsFuseParentProcessDir)
	defer os.Unsetenv(GCSFUSE_PARENT_PROCESS_DIR)

	f := parseArgs(args)

	ExpectEq(filepath.Join(gcsFuseParentProcessDir, "test.txt"), f.LogFile)
	ExpectEq(filepath.Join(gcsFuseParentProcessDir, "test.json"), f.KeyFile)
}
<|endoftext|>


<commit_before>package gaesafedumpserver

import (
	"context"
	"net/http"
	"time"

	"io/ioutil"

	"github.com/continusec/safeadmin"
	"github.com/continusec/safeadmin/pb"
	"github.com/golang/protobuf/proto"
	"google.golang.org/appengine"
	"google.golang.org/appengine/log"
)

func init() {
	server := &safeadmin.SafeDumpServer{
		Storage:                   &GoogleCloudDatastorePersistenceLayer{},
		MaxDecryptionPeriod:       time.Hour * 24 * 7,
		CertificateRotationPeriod: time.Hour * 24,
		PurgeOldKeys:              true,
		KeyRetentionPeriod:        time.Duration(0), // don't store for any additional time
	}

	http.HandleFunc("/simpleRPC/GetPublicCert", func(w http.ResponseWriter, r *http.Request) {
		ctx := appengine.NewContext(r)
		req := &pb.GetPublicCertRequest{}
		if commonStart(ctx, w, r, req) {
			resp, err := server.GetPublicCert(ctx, req)
			commonEnd(ctx, w, resp, err)
		}
	})
	http.HandleFunc("/simpleRPC/DecryptSecret", func(w http.ResponseWriter, r *http.Request) {
		ctx := appengine.NewContext(r)
		req := &pb.DecryptSecretRequest{}
		if commonStart(ctx, w, r, req) {
			resp, err := server.DecryptSecret(ctx, req)
			commonEnd(ctx, w, resp, err)
		}
	})
	http.HandleFunc("/tasks/CronPurge", func(w http.ResponseWriter, r *http.Request) {
		ctx := appengine.NewContext(r)
		err := server.CronPurge(ctx, time.Now())
		if err != nil {
			log.Errorf(ctx, "Error running cron: %s", err)
			w.WriteHeader(http.StatusInternalServerError)
		}
	})
}

// commonStart reads the POST data, and unmarshals into the given message
// Returns True if ready to go. If error occurs, error message is sent
// and returns False
func commonStart(ctx context.Context, w http.ResponseWriter, r *http.Request, req proto.Message) bool {
	if r.Method != http.MethodPost {
		w.WriteHeader(http.StatusBadRequest)
		return false
	}
	defer r.Body.Close()
	input, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.Errorf(ctx, "Error reading request: %s", err)
		w.WriteHeader(http.StatusBadRequest)
		return false
	}

	err = proto.Unmarshal(input, req)
	if err != nil {
		log.Errorf(ctx, "Error unmarshaling: %s", err)
		w.WriteHeader(http.StatusBadRequest)
		return false
	}

	return true
}

// commonEnd marshals the proto, and writes it out, if the given error is clean
func commonEnd(ctx context.Context, w http.ResponseWriter, resp proto.Message, err error) {
	switch err {
	case nil:
	// all good
	case safeadmin.ErrInvalidDate:
		// Send different response on this, since it's a "normal" error
		w.WriteHeader(http.StatusBadRequest)
		return
	default:
		log.Errorf(ctx, "Error decrypting: %s", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	output, err := proto.Marshal(resp)
	if err != nil {
		log.Errorf(ctx, "Error marshaling: %s", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	w.Write(output) // ignore error, too late anyway
}
<commit_msg>Fix incorrect context object<commit_after>package gaesafedumpserver

import (
	"net/http"
	"time"

	"golang.org/x/net/context"

	"io/ioutil"

	"github.com/continusec/safeadmin"
	"github.com/continusec/safeadmin/pb"
	"github.com/golang/protobuf/proto"
	"google.golang.org/appengine"
	"google.golang.org/appengine/log"
)

func init() {
	server := &safeadmin.SafeDumpServer{
		Storage:                   &GoogleCloudDatastorePersistenceLayer{},
		MaxDecryptionPeriod:       time.Hour * 24 * 7,
		CertificateRotationPeriod: time.Hour * 24,
		PurgeOldKeys:              true,
		KeyRetentionPeriod:        time.Duration(0), // don't store for any additional time
	}

	http.HandleFunc("/simpleRPC/GetPublicCert", func(w http.ResponseWriter, r *http.Request) {
		ctx := appengine.NewContext(r)
		req := &pb.GetPublicCertRequest{}
		if commonStart(ctx, w, r, req) {
			resp, err := server.GetPublicCert(ctx, req)
			commonEnd(ctx, w, resp, err)
		}
	})
	http.HandleFunc("/simpleRPC/DecryptSecret", func(w http.ResponseWriter, r *http.Request) {
		ctx := appengine.NewContext(r)
		req := &pb.DecryptSecretRequest{}
		if commonStart(ctx, w, r, req) {
			resp, err := server.DecryptSecret(ctx, req)
			commonEnd(ctx, w, resp, err)
		}
	})
	http.HandleFunc("/tasks/CronPurge", func(w http.ResponseWriter, r *http.Request) {
		ctx := appengine.NewContext(r)
		err := server.CronPurge(ctx, time.Now())
		if err != nil {
			log.Errorf(ctx, "Error running cron: %s", err)
			w.WriteHeader(http.StatusInternalServerError)
		}
	})
}

// commonStart reads the POST data, and unmarshals into the given message
// Returns True if ready to go. If error occurs, error message is sent
// and returns False
func commonStart(ctx context.Context, w http.ResponseWriter, r *http.Request, req proto.Message) bool {
	if r.Method != http.MethodPost {
		w.WriteHeader(http.StatusBadRequest)
		return false
	}
	defer r.Body.Close()
	input, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.Errorf(ctx, "Error reading request: %s", err)
		w.WriteHeader(http.StatusBadRequest)
		return false
	}

	err = proto.Unmarshal(input, req)
	if err != nil {
		log.Errorf(ctx, "Error unmarshaling: %s", err)
		w.WriteHeader(http.StatusBadRequest)
		return false
	}

	return true
}

// commonEnd marshals the proto, and writes it out, if the given error is clean
func commonEnd(ctx context.Context, w http.ResponseWriter, resp proto.Message, err error) {
	switch err {
	case nil:
	// all good
	case safeadmin.ErrInvalidDate:
		// Send different response on this, since it's a "normal" error
		w.WriteHeader(http.StatusBadRequest)
		return
	default:
		log.Errorf(ctx, "Error decrypting: %s", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	output, err := proto.Marshal(resp)
	if err != nil {
		log.Errorf(ctx, "Error marshaling: %s", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	w.Write(output) // ignore error, too late anyway
}
<|endoftext|>


<commit_before>package chat

import (
	"sync"
	"time"
)

// Channel users can be in a channel and chat with each other
type Channel struct {
	Name     string
	lockName *sync.Mutex

	Users     map[string]bool
	lockUsers *sync.Mutex
}

// NewChannel returns a channel
func NewChannel() *Channel {
	return &Channel{
		Users:     make(map[string]bool),
		lockName:  new(sync.Mutex),
		lockUsers: new(sync.Mutex),
	}
}

// AddUser adds a user to this channel
func (c *Channel) AddUser(nickName string) {
	c.lockUsers.Lock()
	defer c.lockUsers.Unlock()
	c.Users[nickName] = true
}

// RemoveUser removes a user from this server
func (c *Channel) RemoveUser(nickName string) {
	c.lockUsers.Lock()
	defer c.lockUsers.Unlock()
	delete(c.Users, nickName)
}

// GetName gets a channel's name
func (c *Channel) GetName() string {
	c.lockName.Lock()
	defer c.lockName.Unlock()
	return c.Name
}

// SetName sets a channel's name
func (c *Channel) SetName(channelName string) {
	c.lockName.Lock()
	defer c.lockName.Unlock()
	c.Name = channelName
}

// GetUserCount returns the number of connected users to this channel
func (c *Channel) GetUserCount() int {
	c.lockUsers.Lock()
	defer c.lockUsers.Unlock()
	return len(c.Users)
}

// GetUsers returns nicknames of connected users
func (c *Channel) GetUsers() map[string]bool {
	c.lockUsers.Lock()
	defer c.lockUsers.Unlock()
	return cloneNickNames(c.Users)
}

// Broadcast sends a message to every user in a channel
func (c *Channel) Broadcast(chatServer *Server, message string) {
	now := time.Now()
	message = now.Format(time.Kitchen) + `-` + message

	users := c.GetUsers()

	for nickName := range users {
		user, err := chatServer.GetUser(nickName)
		// User may no longer be connected to the chat server
		if err != nil {
			c.RemoveUser(nickName)
			continue
		}
		user.outgoing <- message
	}
}

// getNickNames turns a map of nicknames to slice of strings
// This copy is so that we can unlock faster and avoid race
// Since we are looping again through nickNames next and making another blocking call
func cloneNickNames(users map[string]bool) map[string]bool {
	nickNames := make(map[string]bool)
	for nickName := range users {
		nickNames[nickName] = true
	}
	return nickNames
}
<commit_msg>Dont send timestamp in broadcast, time should be interpreted by consumer, maybe more of a global thing?<commit_after>package chat

import (
	"sync"

	"github.com/spring1843/chat-server/src/shared/logs"
)

// Channel users can be in a channel and chat with each other
type Channel struct {
	Name     string
	lockName *sync.Mutex

	Users     map[string]bool
	lockUsers *sync.Mutex
}

// NewChannel returns a channel
func NewChannel() *Channel {
	return &Channel{
		Users:     make(map[string]bool),
		lockName:  new(sync.Mutex),
		lockUsers: new(sync.Mutex),
	}
}

// AddUser adds a user to this channel
func (c *Channel) AddUser(nickName string) {
	c.lockUsers.Lock()
	defer c.lockUsers.Unlock()
	c.Users[nickName] = true
}

// RemoveUser removes a user from this server
func (c *Channel) RemoveUser(nickName string) {
	c.lockUsers.Lock()
	defer c.lockUsers.Unlock()
	delete(c.Users, nickName)
}

// GetName gets a channel's name
func (c *Channel) GetName() string {
	c.lockName.Lock()
	defer c.lockName.Unlock()
	return c.Name
}

// SetName sets a channel's name
func (c *Channel) SetName(channelName string) {
	c.lockName.Lock()
	defer c.lockName.Unlock()
	c.Name = channelName
}

// GetUserCount returns the number of connected users to this channel
func (c *Channel) GetUserCount() int {
	c.lockUsers.Lock()
	defer c.lockUsers.Unlock()
	return len(c.Users)
}

// GetUsers returns nicknames of connected users
func (c *Channel) GetUsers() map[string]bool {
	c.lockUsers.Lock()
	defer c.lockUsers.Unlock()
	return cloneNickNames(c.Users)
}

// Broadcast sends a message to every user in a channel
func (c *Channel) Broadcast(chatServer *Server, message string) {
	users := c.GetUsers()
	for nickName := range users {
		user, err := chatServer.GetUser(nickName)
		// User may no longer be connected to the chat server
		if err != nil {
			c.RemoveUser(nickName)
			logs.Errf(err, "User %s is in channel %s but not on connected to the server", user.GetNickName(), c.GetName())
			continue
		}
		user.SetOutgoing(message)
	}
}

// getNickNames turns a map of nicknames to slice of strings
// This copy is so that we can unlock faster and avoid race
// Since we are looping again through nickNames next and making another blocking call
func cloneNickNames(users map[string]bool) map[string]bool {
	nickNames := make(map[string]bool)
	for nickName := range users {
		nickNames[nickName] = true
	}
	return nickNames
}
<|endoftext|>


<commit_before>package dht

import (
	"testing"

	"github.com/stretchr/testify/require"
)

func TestAnnounceNoStartingNodes(t *testing.T) {
	s, err := NewServer(&ServerConfig{
		Conn:       mustListen(":0"),
		NoSecurity: true,
	})
	require.NoError(t, err)
	defer s.Close()
	var ih [20]byte
	copy(ih[:], "blah")
	_, err = s.Announce(ih, 0, true)
	require.EqualError(t, err, "no initial nodes")
}
<commit_msg>Add test logging bloom filter characteristics<commit_after>package dht

import (
	"testing"

	"github.com/stretchr/testify/require"
)

func TestAnnounceNoStartingNodes(t *testing.T) {
	s, err := NewServer(&ServerConfig{
		Conn:       mustListen(":0"),
		NoSecurity: true,
	})
	require.NoError(t, err)
	defer s.Close()
	var ih [20]byte
	copy(ih[:], "blah")
	_, err = s.Announce(ih, 0, true)
	require.EqualError(t, err, "no initial nodes")
}

func TestDefaultTraversalBloomFilterCharacteristics(t *testing.T) {
	bf := newBloomFilterForTraversal()
	t.Logf("%d bits with %d hashes per item", bf.Cap(), bf.K())
}
<|endoftext|>


<commit_before>package ipctransitrouter

import (
	"github.com/kr/pretty"
	"github.com/stretchr/testify/assert"
	"testing"
)

//See tests of the reference implementation:
//https://github.com/dana/perl-IPC-Transit-Router/blob/master/t/basic.t
//https://github.com/dana/perl-Message-Router/blob/master/t/basic.t
//https://github.com/dana/perl-Message-Router/blob/master/t/hash_of_routes.t
//func Route(sendMessage map[string]interface{}, config map[string]interface{}) error {
func TestBasic(t *testing.T) {
	assert := assert.New(t)
	message := map[string]interface{}{
		"a": "b",
	}
	config := map[string]interface{}{
		"routes": []interface{}{"hi", "there"},
	}
	err := Route(message, config)
	assert.Nil(err)
}

func routeTestNo() {
	pretty.Println("Yeah")
}
<commit_msg>more notes<commit_after>package ipctransitrouter

//totally totally look at this:
//https://godoc.org/github.com/stretchr/stew/objects
//up vote
//OR you can use objects.Map from our stew package, it gives you dot accessors for maps:
//
//objects.Map(data).Get("service.auth.token")
//see http://godoc.org/github.com/stretchr/stew/objects
//via http://stackoverflow.com/questions/17056044/golang-quickly-access-data-of-maps-within-maps
import (
	"github.com/kr/pretty"
	"github.com/stretchr/testify/assert"
	"testing"
)

//See tests of the reference implementation:
//https://github.com/dana/perl-IPC-Transit-Router/blob/master/t/basic.t
//https://github.com/dana/perl-Message-Router/blob/master/t/basic.t
//https://github.com/dana/perl-Message-Router/blob/master/t/hash_of_routes.t
//func Route(sendMessage map[string]interface{}, config map[string]interface{}) error {
func TestBasic(t *testing.T) {
	assert := assert.New(t)
	message := map[string]interface{}{
		"a": "b",
	}
	config := map[string]interface{}{
		"routes": []interface{}{"hi", "there"},
	}
	err := Route(message, config)
	assert.Nil(err)
}

func routeTestNo() {
	pretty.Println("Yeah")
}
<|endoftext|>


<commit_before>package collection

import (
	"encoding/json"
	"errors"
	"github.com/kuzzleio/sdk-go/types"
)

// Count returns the number of documents matching the provided set of filters.
// There is a small delay between documents creation and their existence in our advanced search layer,
// usually a couple of seconds.
// That means that a document that was just been created won’t be returned by this function
func (dc Collection) Count(filters types.SearchFilters, options types.QueryOptions) (int, error) {
	type countResult struct {
		Count int `json:"count"`
	}

	ch := make(chan types.KuzzleResponse)
	query := types.KuzzleRequest{
		Collection: dc.collection,
		Index:      dc.index,
		Controller: "document",
		Action:     "count",
		Body:       filters,
	}
	go dc.Kuzzle.Query(query, options, ch)

	res := <-ch

	if res.Error.Message != "" {
		return 0, errors.New(res.Error.Message)
	}
	result := countResult{}
	json.Unmarshal(res.Result, &result)

	return result.Count, nil
}
<commit_msg>Fix tests<commit_after>package collection

import (
	"encoding/json"
	"errors"
	"github.com/kuzzleio/sdk-go/types"
)

// Count returns the number of documents matching the provided set of filters.
// There is a small delay between documents creation and their existence in our advanced search layer,
// usually a couple of seconds.
// That means that a document that was just been created won’t be returned by this function
func (dc Collection) Count(filters *types.SearchFilters, options types.QueryOptions) (int, error) {
	type countResult struct {
		Count int `json:"count"`
	}

	ch := make(chan types.KuzzleResponse)
	query := types.KuzzleRequest{
		Collection: dc.collection,
		Index:      dc.index,
		Controller: "document",
		Action:     "count",
		Body:       filters,
	}
	go dc.Kuzzle.Query(query, options, ch)

	res := <-ch

	if res.Error.Message != "" {
		return 0, errors.New(res.Error.Message)
	}
	result := countResult{}
	json.Unmarshal(res.Result, &result)

	return result.Count, nil
}
<|endoftext|>


<commit_before>package octoyaml

/*
This module handles reading OctoDNS yaml files.  Sadly the YAML files
are so entirely flexible that parsing them is a nighmare.  We UnMarshalYAML
them into a slice of interfaces mapped to interfaces, then use reflection
to walk the tree, interpreting what we find along the way.  As we collect
data we output models.RecordConfig objects.
*/

import (
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"reflect"
	"strconv"

	yaml "gopkg.in/yaml.v2"

	"github.com/StackExchange/dnscontrol/v3/models"
)

// ReadYaml parses a yaml input and returns a list of RecordConfigs
func ReadYaml(r io.Reader, origin string) (models.Records, error) {
	results := models.Records{}

	// Slurp the YAML into a string.
	ydata, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, fmt.Errorf("can not read yaml filehandle: %w", err)
	}

	// Unmarshal the mystery data into a structure we can relect into.
	var mysterydata map[string]interface{}
	err = yaml.Unmarshal(ydata, &mysterydata)
	if err != nil {
		return nil, fmt.Errorf("could not unmarshal yaml: %w", err)
	}
	//fmt.Printf("ReadYaml: mysterydata == %v\n", mysterydata)

	// Traverse every key/value pair.
	for k, v := range mysterydata { // Each label
		// k, v: k is the label, v is everything we know about the label.
		// In other code, k1, v2 refers to one level deeper, k3, k3 refers to
		// one more level deeper, and so on.
		//fmt.Printf("ReadYaml: NEXT KEY\n")
		//fmt.Printf("ReadYaml:  KEY=%s v.(type)=%s\n", k, reflect.TypeOf(v).String())
		switch v.(type) {
		case map[interface{}]interface{}:
			// The value is itself a map. This means we have a label with
			// with one or more records, each of them are all the same rtype.
			// parseLeaf will handle both of these forms:
			// For example, this:
			// 'www':
			//    type: A
			//    values:
			//      - 1.2.3.4
			//      - 1.2.3.5
			// or
			// 'www':
			//    type: CNAME
			//    value: foo.example.com.
			results, err = parseLeaf(results, k, v, origin)
			if err != nil {
				return results, fmt.Errorf("leaf (%v) error: %w", v, err)
			}
		case []interface{}:
			// The value is a list. This means we have a label with
			// multiple records, each of them may be different rtypes.
			// We need to call parseLeaf() once for each rtype.
			// For example, this:
			// 'www':
			// - type: A
			//   values:
			// 	- 1.2.3.4
			// 	- 1.2.3.5
			// - type: MX
			//   values:
			// 	- priority: 10
			// 	  value: mx1.example.com.
			// 	- priority: 10
			// 	  value: mx2.example.com.
			for i, v3 := range v.([]interface{}) { // All the label's list
				_ = i
				//fmt.Printf("ReadYaml:   list key=%s i=%d v3.(type)=%s\n", k, i, typeof(v3))
				switch v3.(type) {
				case map[interface{}]interface{}:
					//fmt.Printf("ReadYaml:   v3=%v\n", v3)
					results, err = parseLeaf(results, k, v3, origin)
					if err != nil {
						return results, fmt.Errorf("leaf v3=%v: %w", v3, err)
					}
				default:
					return nil, fmt.Errorf("unknown type in list3: k=%s v.(type)=%T v=%v", k, v, v)
				}
			}

		default:
			return nil, fmt.Errorf("unknown type in list1: k=%s v.(type)=%T v=%v", k, v, v)
		}
	}

	sortRecs(results, origin)
	//fmt.Printf("ReadYaml: RESULTS=%v\n", results)
	return results, nil
}

func parseLeaf(results models.Records, k string, v interface{}, origin string) (models.Records, error) {
	var rType, rTarget string
	var rTTL uint32
	rTargets := []string{}
	var someresults models.Records
	for k2, v2 := range v.(map[interface{}]interface{}) { // All  the label's items
		// fmt.Printf("ReadYaml: ifs tk2=%s tv2=%s len(rTargets)=%d\n", typeof(k2), typeof(v2), len(rTargets))
		if typeof(k2) == "string" && (typeof(v2) == "string" || typeof(v2) == "int") {
			// The 2nd level key is a string, and the 2nd level value is a string or int.
			// Here are 3 examples:
			// type: CNAME
			// value: foo.example.com.
			// ttl: 3
			//fmt.Printf("parseLeaf:   k2=%s v2=%v\n", k2, v2)
			switch k2.(string) {
			case "type":
				rType = v2.(string)
			case "ttl":
				var err error
				rTTL, err = decodeTTL(v2)
				if err != nil {
					return nil, fmt.Errorf("parseLeaf: can not parse ttl (%v)", v2)
				}
			case "value":
				rTarget = v2.(string)
			case "values":
				switch v2.(type) {
				case string:
					rTarget = v2.(string)
				default:
					return nil, fmt.Errorf("parseLeaf: unknown type in values: rtpe=%s k=%s k2=%s v2.(type)=%T v2=%v", rType, k, k2, v2, v2)
				}
			default:
				panic("Should not happen")
			}
		} else if typeof(k2) == "string" && typeof(v2) == "[]interface {}" {
			// The 2nd level key is a string, and the 2nd level value is a list.
			someresults = nil
			for _, v3 := range v2.([]interface{}) {
				switch v3.(type) {
				case string:
					// Example:
					// values:
					//   - 1.2.3.1
					//   - 1.2.3.2
					//   - 1.2.3.3
					// We collect all the values for later, when we'll need to generate
					// one RecordConfig for each value.
					//fmt.Printf("parseLeaf: s-append %s\n", v3.(string))
					rTargets = append(rTargets, v3.(string))
				case map[interface{}]interface{}:
					// Example:
					// values:
					// - priority: 10
					//   value: mx1.example.com.
					// - priority: 10
					//   value: mx2.example.com.
					// We collect the individual values. When we are done with this level,
					// we should have enough to generate a single RecordConfig.
					newRc := newRecordConfig(k, rType, "", rTTL, origin)
					for k4, v4 := range v3.(map[interface{}]interface{}) {
						//fmt.Printf("parseLeaf: k4=%s v4=%s\n", k4, v4)
						switch k4.(string) {
						case "priority": // MX,SRV
							priority := uint16(v4.(int))
							newRc.MxPreference = priority
							newRc.SrvPriority = priority
							// Assign it to both places. We'll zap the wrong one later.
						case "weight": // SRV
							newRc.SrvWeight = uint16(v4.(int))
						case "port": // SRV
							newRc.SrvPort = uint16(v4.(int))
						case "value": // MX
							newRc.SetTarget(v4.(string))
						}
					}
					//fmt.Printf("parseLeaf: append %v\n", newRc)
					someresults = append(someresults, newRc)
				default:
					return nil, fmt.Errorf("parseLeaf: unknown type in map: rtype=%s k=%s v3.(type)=%T v3=%v", rType, k, v3, v3)
				}
			}
		} else {
			return nil, fmt.Errorf("parseLeaf: unknown type in level 2: k=%s k2=%s v.2(type)=%T v2=%v", k, k2, v2, v2)
		}
	}
	// fmt.Printf("parseLeaf: Target=(%v)\n", rTarget)
	// fmt.Printf("parseLeaf: len(rTargets)=%d\n", len(rTargets))
	// fmt.Printf("parseLeaf: len(someresults)=%d\n", len(someresults))

	// We've now looped through everything about one label. Make the RecordConfig(s).

	if len(someresults) > 0 {
		// We have many results. Generate a RecordConfig for each one.
		for _, r := range someresults {
			r.Type = rType
			r.TTL = rTTL
			results = append(results, r)
			// Earlier we didn't know what the priority was for. Now that  we know the rType,
			// we zap the wrong one.
			switch r.Type {
			case "MX":
				r.SrvPriority = 0
			case "SRV":
				r.MxPreference = 0
			default:
				panic("ugh")
			}
		}
	} else if rTarget != "" && len(rTargets) == 0 {
		// The file used "value".  Generate a single RecordConfig
		//fmt.Printf("parseLeaf: 1-newRecordConfig(%v, %v, %v, %v, %v)\n", k, rType, rTarget, rTTL, origin)
		results = append(results, newRecordConfig(k, rType, rTarget, rTTL, origin))
	} else {
		// The file used "values" so now we must generate a RecordConfig for each value.
		for _, target := range rTargets {
			//fmt.Printf("parseLeaf: 3-newRecordConfig(%v, %v, %v, %v, %v)\n", k, rType, target, rTTL, origin)
			results = append(results, newRecordConfig(k, rType, target, rTTL, origin))
		}
	}
	return results, nil
}

// newRecordConfig is a RecordConfig factory.
func newRecordConfig(rname, rtype, target string, ttl uint32, origin string) *models.RecordConfig {
	rc := &models.RecordConfig{
		Type: rtype,
		TTL:  ttl,
	}
	rc.SetLabel(rname, origin)
	switch rtype {
	case "TXT":
		rc.SetTargetTXT(target)
	default:
		rc.SetTarget(target)
	}
	return rc
}

// typeof returns a string that indicates v's type:
func typeof(v interface{}) string {
	// Cite: https://stackoverflow.com/a/20170555/71978
	return reflect.TypeOf(v).String()
}

// decodeTTL decodes an interface into a TTL value.
// This is useful when you don't know if a TTL arrived as a string or int.
func decodeTTL(ttl interface{}) (uint32, error) {
	switch ttl.(type) {
	case uint32:
		return ttl.(uint32), nil
	case string:
		s := ttl.(string)
		t, err := strconv.ParseUint(s, 10, 32)
		return uint32(t), fmt.Errorf("decodeTTL failed to parse (%s): %w", s, err)
	case int:
		i := ttl.(int)
		if i < 0 || i > math.MaxUint32 {
			return 0, fmt.Errorf("ttl won't fit in 32-bits (%d)", i)
		}
		return uint32(i), nil
	}
	panic("I don't know what type this TTL is")
}
<commit_msg>Fix octodns: constant 4294967295 overflows int (Issue #736) (#738)<commit_after>package octoyaml

/*
This module handles reading OctoDNS yaml files.  Sadly the YAML files
are so entirely flexible that parsing them is a nighmare.  We UnMarshalYAML
them into a slice of interfaces mapped to interfaces, then use reflection
to walk the tree, interpreting what we find along the way.  As we collect
data we output models.RecordConfig objects.
*/

import (
	"fmt"
	"io"
	"io/ioutil"
	"reflect"
	"strconv"

	yaml "gopkg.in/yaml.v2"

	"github.com/StackExchange/dnscontrol/v3/models"
)

// ReadYaml parses a yaml input and returns a list of RecordConfigs
func ReadYaml(r io.Reader, origin string) (models.Records, error) {
	results := models.Records{}

	// Slurp the YAML into a string.
	ydata, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, fmt.Errorf("can not read yaml filehandle: %w", err)
	}

	// Unmarshal the mystery data into a structure we can relect into.
	var mysterydata map[string]interface{}
	err = yaml.Unmarshal(ydata, &mysterydata)
	if err != nil {
		return nil, fmt.Errorf("could not unmarshal yaml: %w", err)
	}
	//fmt.Printf("ReadYaml: mysterydata == %v\n", mysterydata)

	// Traverse every key/value pair.
	for k, v := range mysterydata { // Each label
		// k, v: k is the label, v is everything we know about the label.
		// In other code, k1, v2 refers to one level deeper, k3, k3 refers to
		// one more level deeper, and so on.
		//fmt.Printf("ReadYaml: NEXT KEY\n")
		//fmt.Printf("ReadYaml:  KEY=%s v.(type)=%s\n", k, reflect.TypeOf(v).String())
		switch v.(type) {
		case map[interface{}]interface{}:
			// The value is itself a map. This means we have a label with
			// with one or more records, each of them are all the same rtype.
			// parseLeaf will handle both of these forms:
			// For example, this:
			// 'www':
			//    type: A
			//    values:
			//      - 1.2.3.4
			//      - 1.2.3.5
			// or
			// 'www':
			//    type: CNAME
			//    value: foo.example.com.
			results, err = parseLeaf(results, k, v, origin)
			if err != nil {
				return results, fmt.Errorf("leaf (%v) error: %w", v, err)
			}
		case []interface{}:
			// The value is a list. This means we have a label with
			// multiple records, each of them may be different rtypes.
			// We need to call parseLeaf() once for each rtype.
			// For example, this:
			// 'www':
			// - type: A
			//   values:
			// 	- 1.2.3.4
			// 	- 1.2.3.5
			// - type: MX
			//   values:
			// 	- priority: 10
			// 	  value: mx1.example.com.
			// 	- priority: 10
			// 	  value: mx2.example.com.
			for i, v3 := range v.([]interface{}) { // All the label's list
				_ = i
				//fmt.Printf("ReadYaml:   list key=%s i=%d v3.(type)=%s\n", k, i, typeof(v3))
				switch v3.(type) {
				case map[interface{}]interface{}:
					//fmt.Printf("ReadYaml:   v3=%v\n", v3)
					results, err = parseLeaf(results, k, v3, origin)
					if err != nil {
						return results, fmt.Errorf("leaf v3=%v: %w", v3, err)
					}
				default:
					return nil, fmt.Errorf("unknown type in list3: k=%s v.(type)=%T v=%v", k, v, v)
				}
			}

		default:
			return nil, fmt.Errorf("unknown type in list1: k=%s v.(type)=%T v=%v", k, v, v)
		}
	}

	sortRecs(results, origin)
	//fmt.Printf("ReadYaml: RESULTS=%v\n", results)
	return results, nil
}

func parseLeaf(results models.Records, k string, v interface{}, origin string) (models.Records, error) {
	var rType, rTarget string
	var rTTL uint32
	rTargets := []string{}
	var someresults models.Records
	for k2, v2 := range v.(map[interface{}]interface{}) { // All  the label's items
		// fmt.Printf("ReadYaml: ifs tk2=%s tv2=%s len(rTargets)=%d\n", typeof(k2), typeof(v2), len(rTargets))
		if typeof(k2) == "string" && (typeof(v2) == "string" || typeof(v2) == "int") {
			// The 2nd level key is a string, and the 2nd level value is a string or int.
			// Here are 3 examples:
			// type: CNAME
			// value: foo.example.com.
			// ttl: 3
			//fmt.Printf("parseLeaf:   k2=%s v2=%v\n", k2, v2)
			switch k2.(string) {
			case "type":
				rType = v2.(string)
			case "ttl":
				var err error
				rTTL, err = decodeTTL(v2)
				if err != nil {
					return nil, fmt.Errorf("parseLeaf: can not parse ttl (%v)", v2)
				}
			case "value":
				rTarget = v2.(string)
			case "values":
				switch v2.(type) {
				case string:
					rTarget = v2.(string)
				default:
					return nil, fmt.Errorf("parseLeaf: unknown type in values: rtpe=%s k=%s k2=%s v2.(type)=%T v2=%v", rType, k, k2, v2, v2)
				}
			default:
				panic("Should not happen")
			}
		} else if typeof(k2) == "string" && typeof(v2) == "[]interface {}" {
			// The 2nd level key is a string, and the 2nd level value is a list.
			someresults = nil
			for _, v3 := range v2.([]interface{}) {
				switch v3.(type) {
				case string:
					// Example:
					// values:
					//   - 1.2.3.1
					//   - 1.2.3.2
					//   - 1.2.3.3
					// We collect all the values for later, when we'll need to generate
					// one RecordConfig for each value.
					//fmt.Printf("parseLeaf: s-append %s\n", v3.(string))
					rTargets = append(rTargets, v3.(string))
				case map[interface{}]interface{}:
					// Example:
					// values:
					// - priority: 10
					//   value: mx1.example.com.
					// - priority: 10
					//   value: mx2.example.com.
					// We collect the individual values. When we are done with this level,
					// we should have enough to generate a single RecordConfig.
					newRc := newRecordConfig(k, rType, "", rTTL, origin)
					for k4, v4 := range v3.(map[interface{}]interface{}) {
						//fmt.Printf("parseLeaf: k4=%s v4=%s\n", k4, v4)
						switch k4.(string) {
						case "priority": // MX,SRV
							priority := uint16(v4.(int))
							newRc.MxPreference = priority
							newRc.SrvPriority = priority
							// Assign it to both places. We'll zap the wrong one later.
						case "weight": // SRV
							newRc.SrvWeight = uint16(v4.(int))
						case "port": // SRV
							newRc.SrvPort = uint16(v4.(int))
						case "value": // MX
							newRc.SetTarget(v4.(string))
						}
					}
					//fmt.Printf("parseLeaf: append %v\n", newRc)
					someresults = append(someresults, newRc)
				default:
					return nil, fmt.Errorf("parseLeaf: unknown type in map: rtype=%s k=%s v3.(type)=%T v3=%v", rType, k, v3, v3)
				}
			}
		} else {
			return nil, fmt.Errorf("parseLeaf: unknown type in level 2: k=%s k2=%s v.2(type)=%T v2=%v", k, k2, v2, v2)
		}
	}
	// fmt.Printf("parseLeaf: Target=(%v)\n", rTarget)
	// fmt.Printf("parseLeaf: len(rTargets)=%d\n", len(rTargets))
	// fmt.Printf("parseLeaf: len(someresults)=%d\n", len(someresults))

	// We've now looped through everything about one label. Make the RecordConfig(s).

	if len(someresults) > 0 {
		// We have many results. Generate a RecordConfig for each one.
		for _, r := range someresults {
			r.Type = rType
			r.TTL = rTTL
			results = append(results, r)
			// Earlier we didn't know what the priority was for. Now that  we know the rType,
			// we zap the wrong one.
			switch r.Type {
			case "MX":
				r.SrvPriority = 0
			case "SRV":
				r.MxPreference = 0
			default:
				panic("ugh")
			}
		}
	} else if rTarget != "" && len(rTargets) == 0 {
		// The file used "value".  Generate a single RecordConfig
		//fmt.Printf("parseLeaf: 1-newRecordConfig(%v, %v, %v, %v, %v)\n", k, rType, rTarget, rTTL, origin)
		results = append(results, newRecordConfig(k, rType, rTarget, rTTL, origin))
	} else {
		// The file used "values" so now we must generate a RecordConfig for each value.
		for _, target := range rTargets {
			//fmt.Printf("parseLeaf: 3-newRecordConfig(%v, %v, %v, %v, %v)\n", k, rType, target, rTTL, origin)
			results = append(results, newRecordConfig(k, rType, target, rTTL, origin))
		}
	}
	return results, nil
}

// newRecordConfig is a RecordConfig factory.
func newRecordConfig(rname, rtype, target string, ttl uint32, origin string) *models.RecordConfig {
	rc := &models.RecordConfig{
		Type: rtype,
		TTL:  ttl,
	}
	rc.SetLabel(rname, origin)
	switch rtype {
	case "TXT":
		rc.SetTargetTXT(target)
	default:
		rc.SetTarget(target)
	}
	return rc
}

// typeof returns a string that indicates v's type:
func typeof(v interface{}) string {
	// Cite: https://stackoverflow.com/a/20170555/71978
	return reflect.TypeOf(v).String()
}

// decodeTTL decodes an interface into a TTL value.
// This is useful when you don't know if a TTL arrived as a string or int.
func decodeTTL(ttl interface{}) (uint32, error) {
	switch ttl.(type) {
	case uint32:
		return ttl.(uint32), nil
	case string:
		s := ttl.(string)
		t, err := strconv.ParseUint(s, 10, 32)
		return uint32(t), fmt.Errorf("decodeTTL failed to parse (%s): %w", s, err)
	case int:
		i := ttl.(int)
		if i < 0 {
			return 0, fmt.Errorf("ttl cannot be negative (%d)", i)
		}
		return uint32(i), nil
	}
	panic("I don't know what type this TTL is")
}
<|endoftext|>


<commit_before>// Copyright 2014 Brave New Software

//    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 archive

import (
	"log"
	"os"
	"strings"
	"time"

	"github.com/getlantern/statshub/bigquery"
	"github.com/getlantern/statshub/statshub"
)

const (
	ARCHIVE_TO_BIGQUERY = "ARCHIVE_TO_BIGQUERY"
)

var (
	shouldArchive = strings.ToLower(os.Getenv(ARCHIVE_TO_BIGQUERY)) == "true"
)

// StartArchiving starts a goroutine that continuously archives data at regular intervals
// based on the archiveInterval constant.
func StartArchiving() {
	if shouldArchive {
		log.Printf("Archiving to BigQuery at %s", bigquery.ProjectId)
		archivePeriodically("fallback", 10*time.Minute)
		archivePeriodically("country", 1*time.Hour)
		//archivePeriodically("host", 1*time.Hour)
		archivePeriodically("user", 24*time.Hour)
	} else {
		log.Printf("%s was not \"true\", not archiving to BigQuery", ARCHIVE_TO_BIGQUERY)
	}
}

func archivePeriodically(dim string, interval time.Duration) {
	go func() {
		for {
			nextInterval := time.Now().Truncate(interval).Add(interval)
			waitTime := nextInterval.Sub(time.Now())
			time.Sleep(waitTime)
			if err := archiveToBigQuery(dim, interval); err != nil {
				log.Printf("Unable to archive dimension %s to BigQuery: %s", dim, err)
			}
		}
	}()
}

func archiveToBigQuery(dim string, interval time.Duration) error {
	if statsByDim, err := statshub.QueryDims([]string{dim}); err != nil {
		return err
	} else {
		for dimName, dimStats := range statsByDim {
			if statsTable, err := NewStatsTable(bigquery.ProjectId, bigquery.DATASET_ID, dimName); err != nil {
				return err
			} else {
				return statsTable.WriteStats(dimStats, time.Now().Truncate(interval))
			}
		}
		return nil
	}
}
<commit_msg>getlantern/lantern#1606 Added archiving of new natty statistics<commit_after>// Copyright 2014 Brave New Software

//    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 archive

import (
	"log"
	"os"
	"strings"
	"time"

	"github.com/getlantern/statshub/bigquery"
	"github.com/getlantern/statshub/statshub"
)

const (
	ARCHIVE_TO_BIGQUERY = "ARCHIVE_TO_BIGQUERY"
)

var (
	shouldArchive = strings.ToLower(os.Getenv(ARCHIVE_TO_BIGQUERY)) == "true"
)

// StartArchiving starts a goroutine that continuously archives data at regular intervals
// based on the archiveInterval constant.
func StartArchiving() {
	if shouldArchive {
		log.Printf("Archiving to BigQuery at %s", bigquery.ProjectId)
		archivePeriodically("fallback", 10*time.Minute)
		archivePeriodically("country", 1*time.Hour)
		//archivePeriodically("host", 1*time.Hour)
		archivePeriodically("user", 24*time.Hour)
		archivePeriodically("answererCountry", 1*time.Hour)
		archivePeriodically("offererAnswererCountries", 1*time.Hour)
		archivePeriodically("operatingSystem", 1*time.Hour)
	} else {
		log.Printf("%s was not \"true\", not archiving to BigQuery", ARCHIVE_TO_BIGQUERY)
	}
}

func archivePeriodically(dim string, interval time.Duration) {
	go func() {
		for {
			nextInterval := time.Now().Truncate(interval).Add(interval)
			waitTime := nextInterval.Sub(time.Now())
			time.Sleep(waitTime)
			if err := archiveToBigQuery(dim, interval); err != nil {
				log.Printf("Unable to archive dimension %s to BigQuery: %s", dim, err)
			}
		}
	}()
}

func archiveToBigQuery(dim string, interval time.Duration) error {
	if statsByDim, err := statshub.QueryDims([]string{dim}); err != nil {
		return err
	} else {
		for dimName, dimStats := range statsByDim {
			if statsTable, err := NewStatsTable(bigquery.ProjectId, bigquery.DATASET_ID, dimName); err != nil {
				return err
			} else {
				return statsTable.WriteStats(dimStats, time.Now().Truncate(interval))
			}
		}
		return nil
	}
}
<|endoftext|>