<commit_before>/*package density interpolates sequences of particle positions onto a density
grid.
*/
package density

import (
	"fmt"
	"log"
	"math"

	"github.com/phil-mansfield/num/rand"
	"github.com/phil-mansfield/gotetra/geom"
	"github.com/phil-mansfield/gotetra/catalog"
)

// Interpolator creates a grid-based density distribution from seqeunces of
// positions.
type Interpolator interface {
	// Interpolate adds the density distribution implied by points to the
	// density grid used by the Interpolator. Particles should all be within
	// the bounds of the bounding grid and points not within the interpolation
	// grid will be ignored.
	Interpolate(mass float64, ids []int64, xs []geom.Vec)
}

// Flag indicates the interpolation scheme that should be used to assign
// densities.
type Flag int

type cic struct {
	g, bg      geom.Grid
	cellWidth  float64
	cellVolume float64
	rhos       []float64
}

type ngp struct {
	g          geom.Grid
	cellWidth  float64
	cellVolume float64
	rhos       []float64
}

type mcarlo struct {
	subIntr Interpolator
	man *catalog.ParticleManager
	totalWidth float64
	countWidth int64
	steps int

	gen *rand.Generator

	// Buffers
	idxBuf geom.TetraIdxs
	tet geom.Tetra
	randBuf []float64
	vecBuf []geom.Vec
}

const (
	CloudInCell Flag = iota
	NearestGridPoint
	MonteCarlo
)

const (
	DefaultMonteCarloSteps = 100
	DefaultMonteCarloGenerator = rand.Tausworthe
)

// Bounds returns a large bounding grid and a smaller interpolation Grid
// which acts as a single subcell of the bounding Grid. cells gives the number
// of cells in the interpolation Grid, and gridWidth gives the number of
// interpolation Grids within the bounding grid [on one side].
func Bounds(cells, gridWidth, gx, gy, gz int) (g, bg *geom.Grid) {
	g = geom.NewGrid(&[3]int{gx * cells, gy * cells, gz * cells}, cells)
	bg = geom.NewGrid(&[3]int{0, 0, 0}, cells*gridWidth)
	return g, bg
}

// NewInterpolator creates an Interpolator instance using the given
// interpolation scheme which adds to the grid rhos. rhos has boundaries
// defined by the Grid g which is embedded in the bounding Grid bg. These two
// grids will almost always be possible to create through a call to Bounds. The
// variable width refers to the interpolation grid, not the bounding grid.
//
// TODO: this initializer is really gross and in some cases requires values that
// don't even get used. Rewrite this later.
func NewInterpolator(
	flag Flag, g, bg *geom.Grid,
	width float64, countWidth int64,
	man *catalog.ParticleManager,
	rhos []float64,
) Interpolator {
	if g.Volume != len(rhos) {
		panic(fmt.Sprintf(
			"Volume of grid, %d, does not equal length of rhos, %d.",
			g.Volume, len(rhos),
		))
	}

	cellWidth := width / float64(g.Width)
	cellVolume := cellWidth * cellWidth * cellWidth

	switch flag {
	case CloudInCell:
		return &cic{*g, *bg, cellWidth, cellVolume, rhos}
	case NearestGridPoint:
		return &ngp{*g, cellWidth, cellVolume, rhos}
	case MonteCarlo:
		return NewMonteCarloInterpolator(
			g, bg, width, countWidth, man, DefaultMonteCarloSteps, rhos,
		)
	}
	panic(fmt.Sprintf("Unknown flag %d", flag))
}

// NewMonteCarloInterpolator creates a new Interpolator using the specified
// number of interpolation points.
func NewMonteCarloInterpolator(
	g, bg *geom.Grid,
	width float64, countWidth int64,
	man *catalog.ParticleManager,
	steps int,
	rhos []float64,
) Interpolator {

	subIntr := NewInterpolator(NearestGridPoint, g, bg, width, countWidth, man, rhos)
	totalWidth := width * float64(bg.Width) / float64(g.Width)

	return &mcarlo{
		subIntr, man, totalWidth, countWidth, steps,
		rand.NewTimeSeed(DefaultMonteCarloGenerator),
		geom.TetraIdxs{}, geom.Tetra{},
		make([]float64, steps * 3), make([]geom.Vec, steps),
	}
}

// Interpolate interpolates a sequence of particles onto a density grid via a
// nearest grid point scheme.
func (intr *ngp) Interpolate(mass float64, ids []int64, xs []geom.Vec) {
	frac := mass / intr.cellVolume
	for _, pt := range xs {
		xp, yp, zp := float64(pt[0]), float64(pt[1]), float64(pt[2])
		xc, yc, zc := cellPoints(xp, yp, zp, intr.cellWidth)
		i, j, k := int(xc), int(yc), int(zc)

		if idx, ok := intr.g.IdxCheck(i, j, k); ok {
			intr.rhos[idx] += frac
		}
	}
}

// Interpolate interpolates a sequence of particles onto a density grid via a
// cloud in cell scheme.
func (intr *cic) Interpolate(mass float64, ids []int64, xs []geom.Vec) {
	frac := mass / intr.cellVolume
	cw, cw2 := intr.cellWidth, intr.cellWidth / 2
	for _, pt := range xs {

		xp, yp, zp := float64(pt[0])-cw2, float64(pt[1])-cw2, float64(pt[2])-cw2
		xc, yc, zc := cellPoints(xp, yp, zp, intr.cellWidth)
		dx, dy, dz := (xp / cw)-xc, (yp / cw)-yc, (zp / cw)-zc
		tx, ty, tz := 1-dx, 1-dy, 1-dz

		i0, i1 := intr.nbrs(int(xc))
		j0, j1 := intr.nbrs(int(yc))
		k0, k1 := intr.nbrs(int(zc))

		intr.incr(i0, j0, k0, tx*ty*tz*frac)
		intr.incr(i1, j0, k0, dx*ty*tz*frac)
		intr.incr(i0, j1, k0, tx*dy*tz*frac)
		intr.incr(i1, j1, k0, dx*dy*tz*frac)
		intr.incr(i0, j0, k1, tx*ty*dz*frac)
		intr.incr(i1, j0, k1, dx*ty*dz*frac)
		intr.incr(i0, j1, k1, tx*dy*dz*frac)
		intr.incr(i1, j1, k1, dx*dy*dz*frac)
	}
}

func (intr *cic) nbrs(i int) (i0, i1 int) {
	if i == -1 {
		return intr.bg.Width - 1, 0
	}
	if i+1 == intr.bg.Width {
		return i, 0
	}
	return i, i + 1
}

func (intr *cic) incr(i, j, k int, frac float64) {
	if idx, ok := intr.g.IdxCheck(i, j, k); ok {
		intr.rhos[idx] += frac
	}
}

func cellPoints(x, y, z, cw float64) (xc, yc, zc float64) {
	return math.Floor(x / cw), math.Floor(y / cw), math.Floor(z / cw)
}

func (intr *mcarlo) Interpolate(mass float64, ids []int64, xs []geom.Vec) {
	ptMass := mass / float64(intr.steps) / 6.0

	for _, id := range ids {
		for dir := 0; dir < 6; dir++ {
			intr.idxBuf.Init(id, intr.countWidth, dir)

			p0 := intr.man.Get(intr.idxBuf[0])
			p1 := intr.man.Get(intr.idxBuf[1])
			p2 := intr.man.Get(intr.idxBuf[2])
			p3 := intr.man.Get(intr.idxBuf[3])

			if p0 == nil || p1 == nil || p2 == nil || p3 == nil {
				log.Printf("Tetrahedron [%v %v %v %v] not in manager.\n",
					p0, p1, p2, p3)

				if p0 == nil {
					log.Printf("Point at index %d not in manager\n",
						intr.idxBuf[0])
				} else if p1 == nil {
					log.Printf("Point at index %d not in manager\n",
						intr.idxBuf[1])
				} else if p2 == nil {
					log.Printf("Point at index %d not in manager\n",
						intr.idxBuf[2])
				} else if p3 == nil {
					log.Printf("Point at index %d not in manager\n",
						intr.idxBuf[3])
				}

				continue
			}

			intr.tet.Init(&p0.Xs, &p1.Xs, &p2.Xs, &p3.Xs, intr.totalWidth)
			intr.tet.Sample(intr.gen, intr.randBuf, intr.vecBuf)
			intr.subIntr.Interpolate(ptMass, nil, intr.vecBuf)
		}
	}
}
<commit_msg>Commited GridCenters implementation.<commit_after>/*package density interpolates sequences of particle positions onto a density
grid.
*/
package density

import (
	"log"
	"math"

	"github.com/phil-mansfield/num/rand"
	"github.com/phil-mansfield/gotetra/geom"
	"github.com/phil-mansfield/gotetra/catalog"
)

// Interpolator creates a grid-based density distribution from seqeunces of
// positions.
type Interpolator interface {
	// Interpolate adds the density distribution implied by points to the
	// density grid used by the Interpolator. Particles should all be within
	// the bounds of the bounding grid and points not within the interpolation
	// grid will be ignored.
	Interpolate(gs []Grid, mass float64, ids []int64, xs []geom.Vec)
}

type Grid struct {
	Rhos []float64
	BoxWidth, CellWidth, CellVolume float64
	G, BG geom.Grid
}

type Cell struct {
	Width, X, Y, Z int
}

type cic struct { }
type ngp struct { }

type mcarlo struct {
	subIntr Interpolator
	man *catalog.ParticleManager
	countWidth int64
	steps int

	gen *rand.Generator

	// Buffers
	idxBuf geom.TetraIdxs
	tet geom.Tetra
	randBuf []float64
	vecBuf []geom.Vec
}

type cellCenter struct {
	man *catalog.ParticleManager
	countWidth int64
	idxBuf geom.TetraIdxs
	tet geom.Tetra
}

func NewGrid(boxWidth float64, gridWidth int, rhos []float64, c *Cell) *Grid {
	if len(rhos) != c.Width * c.Width * c.Width {
		panic("Length of rhos doesn't match cell width.")
	}

	g := &Grid{}
	g.G.Init(&[3]int{c.X*c.Width, c.Y*c.Width, c.Z*c.Width}, c.Width)
	g.BG.Init(&[3]int{0, 0, 0}, c.Width * gridWidth)
	g.BoxWidth = boxWidth
	g.CellWidth = boxWidth / float64(c.Width * gridWidth)
	g.CellVolume = g.CellWidth * g.CellWidth * g.CellWidth
	g.Rhos = rhos
	return g
}

func CloudInCell() Interpolator {
	return &cic{}
}

func NearestGridPoint() Interpolator {
	return &ngp{}
}

func CellCenter(man *catalog.ParticleManager, countWidth int64) Interpolator {
	return &cellCenter{man, countWidth, geom.TetraIdxs{}, geom.Tetra{}}
}

func MonteCarlo(man *catalog.ParticleManager, countWidth int64,
	gen *rand.Generator, steps int) Interpolator {

	return &mcarlo{
		CloudInCell(), man, countWidth, steps,
		gen, geom.TetraIdxs{}, geom.Tetra{},
		make([]float64, steps * 3), make([]geom.Vec, steps),
	}
}

func SubMonteCarlo( gs []Grid, countWidth int, steps int) Interpolator {
	panic("Not Yet Implemneted.")
}

// Interpolate interpolates a sequence of particles onto a density grid via a
// nearest grid point scheme.
func (intr *ngp) Interpolate(gs []Grid, mass float64, ids []int64, xs []geom.Vec) {
	frac := mass / gs[0].CellVolume
	for _, pt := range xs {
		xp, yp, zp := float64(pt[0]), float64(pt[1]), float64(pt[2])
		xc, yc, zc := cellPoints(xp, yp, zp, gs[0].CellWidth)
		i, j, k := int(xc), int(yc), int(zc)

		for gIdx := range gs {
			g := &gs[gIdx]
			if idx, ok := g.G.IdxCheck(i, j, k); ok {
				g.Rhos[idx] += frac
				continue
			}
		}
	}
}

// Interpolate interpolates a sequence of particles onto a density grid via a
// cloud in cell scheme.
func (intr *cic) Interpolate(gs []Grid, mass float64, ids []int64, xs []geom.Vec) {
	frac := mass / gs[0].CellVolume
	cw, cw2 := gs[0].CellWidth, gs[0].CellWidth / 2
	for _, pt := range xs {

		xp, yp, zp := float64(pt[0])-cw2, float64(pt[1])-cw2, float64(pt[2])-cw2
		xc, yc, zc := cellPoints(xp, yp, zp, gs[0].CellWidth)
		dx, dy, dz := (xp / cw)-xc, (yp / cw)-yc, (zp / cw)-zc
		tx, ty, tz := 1-dx, 1-dy, 1-dz

		i0, i1 := gs[0].nbrs(int(xc))
		j0, j1 := gs[0].nbrs(int(yc))
		k0, k1 := gs[0].nbrs(int(zc))

		over000 := tx*ty*tz*frac
		over100 := dx*ty*tz*frac
		over010 := tx*dy*tz*frac
		over110 := dx*dy*tz*frac
		over001 := tx*ty*dz*frac
		over101 := dx*ty*dz*frac
		over011 := tx*dy*dz*frac
		over111 := dx*dy*dz*frac

		for gIdx := range gs {
			g := &gs[gIdx]

			g.incr(i0, j0, k0, over000)
			g.incr(i1, j0, k0, over100)
			g.incr(i0, j1, k0, over010)
			g.incr(i1, j1, k0, over110)
			g.incr(i0, j0, k1, over001)
			g.incr(i1, j0, k1, over101)
			g.incr(i0, j1, k1, over011)
			g.incr(i1, j1, k1, over111)
		}
	}
}

func (g *Grid) nbrs(i int) (i0, i1 int) {
	if i == -1 {
		return g.BG.Width - 1, 0
	}
	if i+1 == g.BG.Width {
		return i, 0
	}
	return i, i + 1
}

func (g *Grid) incr(i, j, k int, frac float64) {
	if idx, ok := g.G.IdxCheck(i, j, k); ok {
		g.Rhos[idx] += frac
	}
}

func cellPoints(x, y, z, cw float64) (xc, yc, zc float64) {
	return math.Floor(x / cw), math.Floor(y / cw), math.Floor(z / cw)
}

func (intr *cellCenter) Interpolate(gs []Grid, mass float64, ids []int64, xs []geom.Vec) {
	cb := &geom.CellBounds{}

	for _, id := range ids {
		for dir := 0; dir < 5; dir++ {
			intr.idxBuf.Init(id, intr.countWidth, dir)

			p0 := intr.man.Get(intr.idxBuf[0])
			p1 := intr.man.Get(intr.idxBuf[1])
			p2 := intr.man.Get(intr.idxBuf[2])
			p3 := intr.man.Get(intr.idxBuf[3])

			if p0 == nil || p1 == nil || p2 == nil || p3 == nil {
				log.Printf("Tetrahedron [%v %v %v %v] not in manager.\n",
					p0, p1, p2, p3)

				if p0 == nil {
					log.Printf("Point at index %d not in manager\n",
						intr.idxBuf[0])
				} else if p1 == nil {
					log.Printf("Point at index %d not in manager\n",
						intr.idxBuf[1])
				} else if p2 == nil {
					log.Printf("Point at index %d not in manager\n",
						intr.idxBuf[2])
				} else if p3 == nil {
					log.Printf("Point at index %d not in manager\n",
						intr.idxBuf[3])
				}

				continue
			}

			intr.tet.Init(&p0.Xs, &p1.Xs, &p2.Xs, &p3.Xs, gs[0].BoxWidth)
			intr.tet.CellBoundsAt(gs[0].CellWidth, cb)

			for i := range gs {
				if gs[i].G.Intersect(cb, &gs[i].BG) {
					intr.intrTetra(mass / 6.0, &gs[i], cb)
				}
			}
		}
	}
}

func (intr *cellCenter) intrTetra(mass float64, g *Grid, cb *geom.CellBounds) {
	minX := maxInt(cb.Min[0], g.BG.Origin[0])
	maxX := minInt(cb.Max[0], g.BG.Origin[0] + g.BG.Width)
	minY := maxInt(cb.Min[1], g.BG.Origin[1])
	maxY := minInt(cb.Max[1], g.BG.Origin[1] + g.BG.Width)
	minZ := maxInt(cb.Min[2], g.BG.Origin[2])
	maxZ := minInt(cb.Max[2], g.BG.Origin[2] + g.BG.Width)

	frac := mass * g.CellVolume / intr.tet.Volume()

	for x := minX; x <= maxX; x++ {
		for y := minY; y <= maxY; y++ {
			for z := minZ; z <= maxZ; z++ {
				xIdx, yIdx, zIdx := g.BG.Wrap(x, y, z)
				idx := g.BG.Idx(xIdx, yIdx, zIdx)
				g.Rhos[idx] += frac
			}
		}
	}
}

func maxInt(x, y int) int {
	if x > y {
		return x
	}
	return y
}

func minInt(x, y int) int {
	if x < y {
		return x
	}
	return y
}

func (intr *mcarlo) Interpolate(gs []Grid, mass float64, ids []int64, xs []geom.Vec) {
	ptMass := mass / float64(intr.steps) / 6.0
	intersectGs := make([]Grid, len(gs))
	cb := &geom.CellBounds{}

	for _, id := range ids {
		for dir := 0; dir < 6; dir++ {
			intr.idxBuf.Init(id, intr.countWidth, dir)

			p0 := intr.man.Get(intr.idxBuf[0])
			p1 := intr.man.Get(intr.idxBuf[1])
			p2 := intr.man.Get(intr.idxBuf[2])
			p3 := intr.man.Get(intr.idxBuf[3])

			if p0 == nil || p1 == nil || p2 == nil || p3 == nil {
				log.Printf("Tetrahedron [%v %v %v %v] not in manager.\n",
					p0, p1, p2, p3)

				if p0 == nil {
					log.Printf("Point at index %d not in manager\n",
						intr.idxBuf[0])
				} else if p1 == nil {
					log.Printf("Point at index %d not in manager\n",
						intr.idxBuf[1])
				} else if p2 == nil {
					log.Printf("Point at index %d not in manager\n",
						intr.idxBuf[2])
				} else if p3 == nil {
					log.Printf("Point at index %d not in manager\n",
						intr.idxBuf[3])
				}

				continue
			}

			intr.tet.Init(&p0.Xs, &p1.Xs, &p2.Xs, &p3.Xs, gs[0].BoxWidth)
			intr.tet.Sample(intr.gen, intr.randBuf, intr.vecBuf)
			intr.tet.CellBoundsAt(gs[0].CellWidth, cb)

			intersectNum := 0
			for i := range gs {
				if gs[i].G.Intersect(cb, &gs[i].BG) {
					intersectGs[intersectNum] = gs[i]
					intersectNum++
				}
			}

			intr.subIntr.Interpolate(intersectGs[0: intersectNum],
				ptMass, nil, intr.vecBuf)
		}
	}
}
<|endoftext|>


<commit_before>package messagehandle

import (
	"github.com/canopener/PongPlusPlus-Server/server/connection"
	"github.com/canopener/PongPlusPlus-Server/server/games"
	"github.com/canopener/PongPlusPlus-Server/server/messages"
	"github.com/canopener/serverlog"
)

// CreateGame handles the case where a player sends a CreateGame message
func CreateGame(conn *connection.Connection, allGames map[string]*games.Game, message messages.RequestGameListMessage) {
	serverlog.General("Received RequestGameList message from conn:", conn.Alias)

	if !conn.Registered {
		serverlog.General("Unregistered conn:", conn.Alias, "called RequestGameList")
		denied := messages.NewCreateGameDeniedMessage(message.GameName, "You are not registered")
		conn.Write(denied.Bytes())
		return
	}

	gameList := messages.NewGameListMessage(allGames)
	conn.Write(gameList.Bytes())
}
<commit_msg>Name change<commit_after>package messagehandle

import (
	"github.com/canopener/PongPlusPlus-Server/server/connection"
	"github.com/canopener/PongPlusPlus-Server/server/games"
	"github.com/canopener/PongPlusPlus-Server/server/messages"
	"github.com/canopener/serverlog"
)

// RequestGameList handles the case where a player sends a RequestGameList message
func RequestGameList(conn *connection.Connection, allGames map[string]*games.Game, message messages.RequestGameListMessage) {
	serverlog.General("Received RequestGameList message from conn:", conn.Alias)

	if !conn.Registered {
		serverlog.General("Unregistered conn:", conn.Alias, "called RequestGameList")
		denied := messages.NewCreateGameDeniedMessage(message.GameName, "You are not registered")
		conn.Write(denied.Bytes())
		return
	}

	gameList := messages.NewGameListMessage(allGames)
	conn.Write(gameList.Bytes())
}
<|endoftext|>


<commit_before>package hostsfile

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"testing"
)

func createTestHostsFile() (*os.File, error) {
	f, err := ioutil.TempFile("testdata", "hosts-")

	if err != nil {
		return nil, err
	}

	contents, _ := ioutil.ReadFile("testdata/hosts")

	f.Write(contents)
	f.Seek(0, 0)

	return f, nil
}

func remove(f *os.File) error {
	return os.Remove(f.Name())
}

func TestAddEntryAddsEntry(t *testing.T) {
	hostname, ip := "dev.dev", "192.168.0.1"
	host := NewHost(ip, hostname)

	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}
	h.AddEntry(host)

	// Reset the offset after AddHostsEntry changes it.
	scanner := bufio.NewScanner(f)

	for scanner.Scan() {
		if strings.Contains(scanner.Text(), "dev.dev") {
			return
		}
	}

	t.Error(fmt.Sprintf("Expected to find %s in testdata/hosts", hostname))
}

func TestRemoveEntryRemovesEntry(t *testing.T) {
	hostname, ip := "dev.dev", "192.168.0.1"
	host := NewHost(ip, hostname)

	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}
	h.RemoveEntry(host)

	scanner := bufio.NewScanner(f)

	for scanner.Scan() {
		if !strings.Contains(scanner.Text(), "dev.dev") {
			return
		}
	}

	t.Error(fmt.Sprintf("Expected to find %s in testdata/hosts", hostname))
}

func TestFindEntryFindsEntry(t *testing.T) {
	hostname, ip := "dev1.dev", "192.168.0.1"
	host := NewHost(ip, hostname)

	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}

	h.AddEntry(host)

	entry := h.FindEntry(hostname)

	if entry == nil {
		t.Error(fmt.Sprintf("Unable to find host entry %s.", hostname))
	}
}

func TestListEntriesReturnsEntries(t *testing.T) {
	hostname1, ip1 := "dev1.dev", "192.168.0.1"
	host1 := NewHost(ip1, hostname1)

	hostname2, ip2 := "dev2.dev", "192.168.0.2"
	host2 := NewHost(ip2, hostname2)

	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}

	h.AddEntry(host1)
	h.AddEntry(host2)

	hosts := h.GetEntries()

	if len(hosts) != 2 {
		t.Error(fmt.Sprintf("Expected to find 2 host entry. Found %d instead.", len(hosts)))
	}
}

func TestListEntriesReturnsNoEntries(t *testing.T) {

}

func TestClean(t *testing.T) {
	hostname1, ip1 := "dev1.dev", "192.168.0.1"
	host1 := NewHost(ip1, hostname1)

	hostname2, ip2 := "dev2.dev", "192.168.0.2"
	host2 := NewHost(ip2, hostname2)

	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}

	h.AddEntry(host1)
	h.AddEntry(host2)

	h.Clean()
	hosts := h.GetEntries()

	if len(hosts) != 0 {
		t.Error(fmt.Sprintf("Expected to find 2 host entry. Found %d instead.", len(hosts)))
	}
}
<commit_msg>Increase test coverage for hostsfile.go.<commit_after>package hostsfile

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"testing"
)

func createTestHostsFile() (*os.File, error) {
	f, err := ioutil.TempFile("testdata", "hosts-")
	defer f.Seek(0, 0)

	if err != nil {
		return nil, err
	}

	contents, _ := ioutil.ReadFile("testdata/hosts")
	f.Write(contents)

	return f, nil
}

func convertToReadOnlyHostsFile(f *os.File) (*os.File, error) {
	f.Close()
	f, err := os.OpenFile(f.Name(), os.O_RDONLY, 0666)

	if err != nil {
		return nil, err
	}

	return f, err
}

func remove(f *os.File) error {
	return os.Remove(f.Name())
}

func TestAddEntryAddsEntry(t *testing.T) {
	hostname, ip := "dev.dev", "192.168.0.1"
	host := NewHost(ip, hostname)

	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}
	h.AddEntry(host)

	// Reset the offset after AddHostsEntry changes it.
	scanner := bufio.NewScanner(f)

	for scanner.Scan() {
		if strings.Contains(scanner.Text(), "dev.dev") {
			return // Return as soon as we match the hostname.
		}
	}

	t.Error(fmt.Sprintf("Expected to find %s in testdata/hosts", hostname))
}

func TestAddEntryReturnsError(t *testing.T) {
	hostname, ip := "dev.dev", "192.168.0.1"
	host := NewHost(ip, hostname)

	f, _ := createTestHostsFile()
	f, _ = convertToReadOnlyHostsFile(f)
	defer remove(f)

	h := Hostsfile{f.Name(), f}
	err := h.AddEntry(host)

	if err == nil {
		t.Error("Expected an error to be returned from AddEntry.")
	}
}

func TestRemoveEntryRemovesEntry(t *testing.T) {
	hostname, ip := "dev.dev", "192.168.0.1"
	host := NewHost(ip, hostname)

	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}
	h.AddEntry(host)
	h.RemoveEntry(host)

	scanner := bufio.NewScanner(f)

	for scanner.Scan() {
		if !strings.Contains(scanner.Text(), "dev.dev") {
			return
		}
	}

	t.Error(fmt.Sprintf("Expected to find %s in testdata/hosts", hostname))
}

func TestRemoveEntryReturnsError(t *testing.T) {
	hostname, ip := "dev.dev", "192.168.0.1"
	host := NewHost(ip, hostname)

	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}
	h.AddEntry(host)

	f, _ = convertToReadOnlyHostsFile(f)
	err := h.RemoveEntry(host)

	if err == nil {
		t.Error("Expected an error to be returned from RemoveEntry.")
	}
}

func TestFindEntryFindsEntry(t *testing.T) {
	hostname, ip := "dev1.dev", "192.168.0.1"
	host := NewHost(ip, hostname)

	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}

	h.AddEntry(host)

	entry := h.FindEntry(hostname)

	if entry == nil {
		t.Error(fmt.Sprintf("Unable to find host entry %s.", hostname))
	}
}

func TestFindEntryDoesNotFindsEntry(t *testing.T) {
	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}

	entry := h.FindEntry("dev.dev")

	if entry != nil {
		t.Error(fmt.Sprintf("Unexpectedly found a host entry."))
	}
}

func TestGetEntriesReturnsEntries(t *testing.T) {
	hostname1, ip1 := "dev1.dev", "192.168.0.1"
	host1 := NewHost(ip1, hostname1)

	hostname2, ip2 := "dev2.dev", "192.168.0.2"
	host2 := NewHost(ip2, hostname2)

	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}

	h.AddEntry(host1)
	h.AddEntry(host2)

	hosts := h.GetEntries()

	if numHosts := len(hosts); numHosts != 2 {
		t.Error(fmt.Sprintf("Expected to find 2 host entries. Found %d instead.", numHosts))
	}
}

func TestGetEntriesReturnsDuplicateEntries(t *testing.T) {
	hostname, ip := "dev.dev", "192.168.0.1"
	host := NewHost(ip, hostname)

	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}

	h.AddEntry(host)
	h.AddEntry(host)
	h.AddEntry(host)

	hosts := h.GetEntries()

	if numHosts := len(hosts); numHosts != 3 {
		t.Error(fmt.Sprintf("Expected to find 3 host entries. Found %d instead.", numHosts))
	}
}

func TestListEntriesReturnsNoEntries(t *testing.T) {
	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}
	hosts := h.GetEntries()

	if numHosts := len(hosts); numHosts != 0 {
		t.Error(fmt.Sprintf("Expected to find 0 host entries. Found %d instead.", numHosts))
	}

}

func TestCleanReturnsAllEntries(t *testing.T) {
	hostname1, ip1 := "dev1.dev", "192.168.0.1"
	host1 := NewHost(ip1, hostname1)

	hostname2, ip2 := "dev2.dev", "192.168.0.2"
	host2 := NewHost(ip2, hostname2)

	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}

	h.AddEntry(host1)
	h.AddEntry(host2)

	h.Clean()
	hosts := h.GetEntries()

	if len(hosts) != 0 {
		t.Error(fmt.Sprintf("Expected to find 2 host entry. Found %d instead.", len(hosts)))
	}
}

func TestCleanWithNoEntries(t *testing.T) {
	f, _ := createTestHostsFile()
	defer remove(f)

	h := Hostsfile{f.Name(), f}

	h.Clean()
	hosts := h.GetEntries()

	if numHosts := len(hosts); numHosts != 0 {
		t.Error(fmt.Sprintf("Expected to find 2 host entry. Found %d instead.", numHosts))
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 dag_test

import (
	cryptorand "crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"sync"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/internal/dag"
	"github.com/jacobsa/comeback/internal/graph"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

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

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

// Compute the the reachability partial order for the DAG defined by the
// supplied edges. Don't include self-reachability.
func reachabilityRelation(
	edges map[string][]string) (r map[string][]string, err error) {
	r = make(map[string][]string)

	var visit func(string)
	visit = func(n string) {
		// Have we already computed the result for n?
		if _, ok := r[n]; ok {
			return
		}

		// The set of things reachable from n is the set of things reachable from
		// its successors, plus those successors.
		set := make(map[string]struct{})
		for _, successor := range edges[n] {
			set[successor] = struct{}{}

			visit(successor)
			for _, m := range r[successor] {
				set[m] = struct{}{}
			}
		}

		for m, _ := range set {
			r[n] = append(r[n], m)
		}
	}

	for n, _ := range edges {
		visit(n)
	}

	return
}

// Return the relation composed of pairs (Y, X) for each pair (X, Y) in the
// input relation.
func invertRelation(r map[string][]string) (inverted map[string][]string) {
	inverted = make(map[string][]string)
	for k, vs := range r {
		for _, v := range vs {
			inverted[v] = append(inverted[v], k)
		}
	}

	// Ensure that everything has an entry, even if empty.
	for k, _ := range r {
		if _, ok := inverted[k]; !ok {
			inverted[k] = nil
		}
	}

	return
}

// Create a rand.Rand seeded with a good source.
func makeRandSource() (src *rand.Rand) {
	// Read a seed from a good source.
	var seed int64
	err := binary.Read(cryptorand.Reader, binary.LittleEndian, &seed)
	if err != nil {
		log.Fatalln(err)
	}

	src = rand.New(rand.NewSource(seed))
	return
}

// A dag.Visitor for string nodes that calls through to a canned function.
type visitor struct {
	F func(context.Context, string) error
}

func (v *visitor) Visit(
	ctx context.Context,
	untyped dag.Node) (err error) {
	err = v.F(ctx, untyped.(string))
	return
}

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

const visitParallelism = 8

type VisitTest struct {
	ctx context.Context
}

var _ SetUpInterface = &VisitTest{}

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

func (t *VisitTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx
}

func (t *VisitTest) call(
	startNodes []string,
	findDependencies func(context.Context, string) ([]string, error),
	visit func(context.Context, string) error) (err error) {
	// Convert to a list of dag.Node.
	var startDAGNodes []dag.Node
	for _, n := range startNodes {
		startDAGNodes = append(startDAGNodes, n)
	}

	// Call.
	err = graph.TraverseDAG(
		t.ctx,
		startDAGNodes,
		&dependencyResolver{F: findDependencies},
		&visitor{F: visit},
		visitParallelism)

	return
}

// Run the graph described by the given edges through TraverseDAG, inserting
// sleeps to attempt to shake out races, and make sure we can't catch it
// violating any of its promises.
func (t *TraverseDAGTest) runTest(edges map[string][]string) {
	// Compute the reachability relation for the DAG, and invert it.
	reachable, err := reachabilityRelation(edges)
	AssertEq(nil, err)
	inverseReachable := invertRelation(reachable)

	// Set up a visit function that ensures that a visit hasn't happened out of
	// order, sleeps for awhile, then marks the node as visited.
	randSrc := makeRandSource()

	var mu sync.Mutex
	visited := make(map[string]struct{}) // GUARDED_BY(mu)

	visit := func(ctx context.Context, n string) (err error) {
		mu.Lock()
		defer mu.Unlock()

		// We should have already finished visiting all of the nodes from which
		// this node is reachable.
		for _, p := range inverseReachable[n] {
			_, ok := visited[p]
			if !ok {
				err = fmt.Errorf("Visited %q before finished visiting %q", n, p)
				return
			}
		}

		// Wait a random amount of time, then mark the node as visited and return.
		mu.Unlock()
		time.Sleep(time.Duration(randSrc.Int63n(int64(20 * time.Millisecond))))
		mu.Lock()

		visited[n] = struct{}{}
		return
	}

	// Call.
	err = t.call(edges, visit)
	AssertEq(nil, err)

	// Make sure everything was visited.
	for n, _ := range edges {
		_, ok := visited[n]
		AssertTrue(ok, "n: %q", n)
	}
}

// Generate a tree with a certain depth, where the number of children for each
// node is random. The root node is "root".
func randomTree(depth int) (edges map[string][]string) {
	edges = make(map[string][]string)
	randSrc := makeRandSource()

	nextID := 0
	nextLevel := []string{"root"}

	for depthI := 0; depthI < depth; depthI++ {
		thisLevel := nextLevel
		nextLevel = nil

		for _, parent := range thisLevel {
			// Ensure that there is an entry, even if there are no children.
			edges[parent] = []string{}

			// Add children.
			numChildren := int(randSrc.Int31n(6))
			for childI := 0; childI < numChildren; childI++ {
				child := fmt.Sprintf("%v", nextID)
				nextID++

				nextLevel = append(nextLevel, child)
				edges[parent] = append(edges[parent], child)
			}
		}
	}

	// Ensure that there is an entry for everything in the last level, even if
	// there are no children.
	for _, n := range nextLevel {
		edges[n] = []string{}
	}

	return
}

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

func (t *TraverseDAGTest) EmptyGraph() {
	edges := map[string][]string{}
	t.runTest(edges)
}

func (t *TraverseDAGTest) SingleNodeConnectedComponents() {
	// Graph structure:
	//
	//     A  B  C  D
	//
	edges := map[string][]string{
		"A": {},
		"B": {},
		"C": {},
		"D": {},
	}

	t.runTest(edges)
}

func (t *TraverseDAGTest) SimpleRootedTree() {
	// Graph structure:
	//
	//        A
	//      / |
	//     B  C
	//      / | \
	//     D  E  F
	//      / |  |
	//     G  H  I
	//           | \
	//           J  K
	//
	edges := map[string][]string{
		"A": {"B", "C"},
		"B": {},
		"C": {"D", "E", "F"},
		"D": {},
		"E": {"G", "H"},
		"F": {"I"},
		"G": {},
		"H": {},
		"I": {"J", "K"},
		"J": {},
		"K": {},
	}

	t.runTest(edges)
}

func (t *TraverseDAGTest) SimpleDAG() {
	// Graph structure:
	//
	//        A
	//      /  \
	//     B    C
	//      \  /|
	//        D |
	//         \|
	//          E
	//
	edges := map[string][]string{
		"A": {"B", "C"},
		"B": {"D"},
		"C": {"D", "E"},
		"D": {"E"},
		"E": {},
	}

	t.runTest(edges)
}

func (t *TraverseDAGTest) MultipleConnectedComponents() {
	// Graph structure:
	//
	//        A       E
	//      /  \      |\
	//     B    C     F G
	//      \  /
	//        D
	//
	edges := map[string][]string{
		"A": {"B", "C"},
		"B": {"D"},
		"C": {"D"},
		"D": {},
		"E": {"F", "G"},
		"F": {},
		"G": {},
	}

	t.runTest(edges)
}

func (t *TraverseDAGTest) RedundantRoots() {
	AssertTrue(false, "TODO")
}

func (t *TraverseDAGTest) Cycle() {
	AssertTrue(false, "TODO")
}

func (t *TraverseDAGTest) LargeRootedTree() {
	const depth = 6
	edges := randomTree(depth)
	t.runTest(edges)
}

func (t *TraverseDAGTest) LargeRootedTree_Inverted() {
	const depth = 6
	edges := invertRelation(randomTree(depth))
	t.runTest(edges)
}

func (t *TraverseDAGTest) DependencyResolverReturnsError() {
	AssertTrue(false, "TODO")
}

func (t *TraverseDAGTest) VisitorReturnsError() {
	AssertGt(traverseDAGParallelism, 1)

	// Graph structure:
	//
	//        A
	//      / |
	//     B  C
	//     |  | \
	//     D  E  F
	//      / |  |
	//     G  H  I
	//           | \
	//           J  K
	//
	edges := map[string][]string{
		"A": {"B", "C"},
		"B": {"D"},
		"C": {"E", "F"},
		"D": {},
		"E": {"G", "H"},
		"F": {"I"},
		"G": {},
		"H": {},
		"I": {"J", "K"},
		"J": {},
		"K": {},
	}

	// Visitor behavior:
	//
	//  *  Record the node visited.
	//
	//  *  For C, wait until told and then return an error.
	//
	//  *  For B:
	//     *   Tell C to proceed.
	//     *   Block until cancelled.
	//     *   Close a channel indicating that the context was cancelled.
	//
	cErr := errors.New("taco")
	bReceived := make(chan struct{})
	bCancelled := make(chan struct{})
	visited := make(chan string, len(edges))

	visit := func(ctx context.Context, n string) (err error) {
		visited <- n

		switch n {
		case "C":
			<-bReceived
			err = cErr

		case "B":
			close(bReceived)

			done := ctx.Done()
			AssertNe(nil, done)
			<-done
			close(bCancelled)
		}

		return
	}

	// Traverse.
	err := t.call(edges, visit)
	ExpectEq(cErr, err)

	// B should have seen cancellation.
	<-bCancelled

	// Nothing descending from B or C should have been visted.
	close(visited)
	var nodes []string
	for n := range visited {
		nodes = append(nodes, n)
	}

	ExpectThat(sortNodes(nodes), ElementsAre("A", "B", "C"))
}
<commit_msg>Removed an old import.<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 dag_test

import (
	cryptorand "crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"sync"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/internal/dag"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

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

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

// Compute the the reachability partial order for the DAG defined by the
// supplied edges. Don't include self-reachability.
func reachabilityRelation(
	edges map[string][]string) (r map[string][]string, err error) {
	r = make(map[string][]string)

	var visit func(string)
	visit = func(n string) {
		// Have we already computed the result for n?
		if _, ok := r[n]; ok {
			return
		}

		// The set of things reachable from n is the set of things reachable from
		// its successors, plus those successors.
		set := make(map[string]struct{})
		for _, successor := range edges[n] {
			set[successor] = struct{}{}

			visit(successor)
			for _, m := range r[successor] {
				set[m] = struct{}{}
			}
		}

		for m, _ := range set {
			r[n] = append(r[n], m)
		}
	}

	for n, _ := range edges {
		visit(n)
	}

	return
}

// Return the relation composed of pairs (Y, X) for each pair (X, Y) in the
// input relation.
func invertRelation(r map[string][]string) (inverted map[string][]string) {
	inverted = make(map[string][]string)
	for k, vs := range r {
		for _, v := range vs {
			inverted[v] = append(inverted[v], k)
		}
	}

	// Ensure that everything has an entry, even if empty.
	for k, _ := range r {
		if _, ok := inverted[k]; !ok {
			inverted[k] = nil
		}
	}

	return
}

// Create a rand.Rand seeded with a good source.
func makeRandSource() (src *rand.Rand) {
	// Read a seed from a good source.
	var seed int64
	err := binary.Read(cryptorand.Reader, binary.LittleEndian, &seed)
	if err != nil {
		log.Fatalln(err)
	}

	src = rand.New(rand.NewSource(seed))
	return
}

// A dag.Visitor for string nodes that calls through to a canned function.
type visitor struct {
	F func(context.Context, string) error
}

func (v *visitor) Visit(
	ctx context.Context,
	untyped dag.Node) (err error) {
	err = v.F(ctx, untyped.(string))
	return
}

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

const visitParallelism = 8

type VisitTest struct {
	ctx context.Context
}

var _ SetUpInterface = &VisitTest{}

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

func (t *VisitTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx
}

func (t *VisitTest) call(
	startNodes []string,
	findDependencies func(context.Context, string) ([]string, error),
	visit func(context.Context, string) error) (err error) {
	// Convert to a list of dag.Node.
	var startDAGNodes []dag.Node
	for _, n := range startNodes {
		startDAGNodes = append(startDAGNodes, n)
	}

	// Call.
	err = graph.TraverseDAG(
		t.ctx,
		startDAGNodes,
		&dependencyResolver{F: findDependencies},
		&visitor{F: visit},
		visitParallelism)

	return
}

// Run the graph described by the given edges through TraverseDAG, inserting
// sleeps to attempt to shake out races, and make sure we can't catch it
// violating any of its promises.
func (t *TraverseDAGTest) runTest(edges map[string][]string) {
	// Compute the reachability relation for the DAG, and invert it.
	reachable, err := reachabilityRelation(edges)
	AssertEq(nil, err)
	inverseReachable := invertRelation(reachable)

	// Set up a visit function that ensures that a visit hasn't happened out of
	// order, sleeps for awhile, then marks the node as visited.
	randSrc := makeRandSource()

	var mu sync.Mutex
	visited := make(map[string]struct{}) // GUARDED_BY(mu)

	visit := func(ctx context.Context, n string) (err error) {
		mu.Lock()
		defer mu.Unlock()

		// We should have already finished visiting all of the nodes from which
		// this node is reachable.
		for _, p := range inverseReachable[n] {
			_, ok := visited[p]
			if !ok {
				err = fmt.Errorf("Visited %q before finished visiting %q", n, p)
				return
			}
		}

		// Wait a random amount of time, then mark the node as visited and return.
		mu.Unlock()
		time.Sleep(time.Duration(randSrc.Int63n(int64(20 * time.Millisecond))))
		mu.Lock()

		visited[n] = struct{}{}
		return
	}

	// Call.
	err = t.call(edges, visit)
	AssertEq(nil, err)

	// Make sure everything was visited.
	for n, _ := range edges {
		_, ok := visited[n]
		AssertTrue(ok, "n: %q", n)
	}
}

// Generate a tree with a certain depth, where the number of children for each
// node is random. The root node is "root".
func randomTree(depth int) (edges map[string][]string) {
	edges = make(map[string][]string)
	randSrc := makeRandSource()

	nextID := 0
	nextLevel := []string{"root"}

	for depthI := 0; depthI < depth; depthI++ {
		thisLevel := nextLevel
		nextLevel = nil

		for _, parent := range thisLevel {
			// Ensure that there is an entry, even if there are no children.
			edges[parent] = []string{}

			// Add children.
			numChildren := int(randSrc.Int31n(6))
			for childI := 0; childI < numChildren; childI++ {
				child := fmt.Sprintf("%v", nextID)
				nextID++

				nextLevel = append(nextLevel, child)
				edges[parent] = append(edges[parent], child)
			}
		}
	}

	// Ensure that there is an entry for everything in the last level, even if
	// there are no children.
	for _, n := range nextLevel {
		edges[n] = []string{}
	}

	return
}

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

func (t *TraverseDAGTest) EmptyGraph() {
	edges := map[string][]string{}
	t.runTest(edges)
}

func (t *TraverseDAGTest) SingleNodeConnectedComponents() {
	// Graph structure:
	//
	//     A  B  C  D
	//
	edges := map[string][]string{
		"A": {},
		"B": {},
		"C": {},
		"D": {},
	}

	t.runTest(edges)
}

func (t *TraverseDAGTest) SimpleRootedTree() {
	// Graph structure:
	//
	//        A
	//      / |
	//     B  C
	//      / | \
	//     D  E  F
	//      / |  |
	//     G  H  I
	//           | \
	//           J  K
	//
	edges := map[string][]string{
		"A": {"B", "C"},
		"B": {},
		"C": {"D", "E", "F"},
		"D": {},
		"E": {"G", "H"},
		"F": {"I"},
		"G": {},
		"H": {},
		"I": {"J", "K"},
		"J": {},
		"K": {},
	}

	t.runTest(edges)
}

func (t *TraverseDAGTest) SimpleDAG() {
	// Graph structure:
	//
	//        A
	//      /  \
	//     B    C
	//      \  /|
	//        D |
	//         \|
	//          E
	//
	edges := map[string][]string{
		"A": {"B", "C"},
		"B": {"D"},
		"C": {"D", "E"},
		"D": {"E"},
		"E": {},
	}

	t.runTest(edges)
}

func (t *TraverseDAGTest) MultipleConnectedComponents() {
	// Graph structure:
	//
	//        A       E
	//      /  \      |\
	//     B    C     F G
	//      \  /
	//        D
	//
	edges := map[string][]string{
		"A": {"B", "C"},
		"B": {"D"},
		"C": {"D"},
		"D": {},
		"E": {"F", "G"},
		"F": {},
		"G": {},
	}

	t.runTest(edges)
}

func (t *TraverseDAGTest) RedundantRoots() {
	AssertTrue(false, "TODO")
}

func (t *TraverseDAGTest) Cycle() {
	AssertTrue(false, "TODO")
}

func (t *TraverseDAGTest) LargeRootedTree() {
	const depth = 6
	edges := randomTree(depth)
	t.runTest(edges)
}

func (t *TraverseDAGTest) LargeRootedTree_Inverted() {
	const depth = 6
	edges := invertRelation(randomTree(depth))
	t.runTest(edges)
}

func (t *TraverseDAGTest) DependencyResolverReturnsError() {
	AssertTrue(false, "TODO")
}

func (t *TraverseDAGTest) VisitorReturnsError() {
	AssertGt(traverseDAGParallelism, 1)

	// Graph structure:
	//
	//        A
	//      / |
	//     B  C
	//     |  | \
	//     D  E  F
	//      / |  |
	//     G  H  I
	//           | \
	//           J  K
	//
	edges := map[string][]string{
		"A": {"B", "C"},
		"B": {"D"},
		"C": {"E", "F"},
		"D": {},
		"E": {"G", "H"},
		"F": {"I"},
		"G": {},
		"H": {},
		"I": {"J", "K"},
		"J": {},
		"K": {},
	}

	// Visitor behavior:
	//
	//  *  Record the node visited.
	//
	//  *  For C, wait until told and then return an error.
	//
	//  *  For B:
	//     *   Tell C to proceed.
	//     *   Block until cancelled.
	//     *   Close a channel indicating that the context was cancelled.
	//
	cErr := errors.New("taco")
	bReceived := make(chan struct{})
	bCancelled := make(chan struct{})
	visited := make(chan string, len(edges))

	visit := func(ctx context.Context, n string) (err error) {
		visited <- n

		switch n {
		case "C":
			<-bReceived
			err = cErr

		case "B":
			close(bReceived)

			done := ctx.Done()
			AssertNe(nil, done)
			<-done
			close(bCancelled)
		}

		return
	}

	// Traverse.
	err := t.call(edges, visit)
	ExpectEq(cErr, err)

	// B should have seen cancellation.
	<-bCancelled

	// Nothing descending from B or C should have been visted.
	close(visited)
	var nodes []string
	for n := range visited {
		nodes = append(nodes, n)
	}

	ExpectThat(sortNodes(nodes), ElementsAre("A", "B", "C"))
}
<|endoftext|>


<commit_before>package memtable

import "sync"

type rwmutex struct {
	sync.RWMutex
}

func (n *node) Next(i int) *node {
	return n.nexts[i]
}

func (n *node) SetNext(i int, next *node) {
	n.nexts[i] = next
}
<commit_msg>Fix missing build tag<commit_after>// +build !atomic

package memtable

import "sync"

type rwmutex struct {
	sync.RWMutex
}

func (n *node) Next(i int) *node {
	return n.nexts[i]
}

func (n *node) SetNext(i int, next *node) {
	n.nexts[i] = next
}
<|endoftext|>


<commit_before>// Package s3 provides a Pipe that push artifacts to s3/minio
package s3

import (
	"os"
	"path/filepath"

	"github.com/apex/log"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/goreleaser/goreleaser/internal/artifact"
	"github.com/goreleaser/goreleaser/internal/pipeline"
	"github.com/goreleaser/goreleaser/internal/semerrgroup"
	"github.com/goreleaser/goreleaser/internal/tmpl"
	"github.com/goreleaser/goreleaser/pkg/config"
	"github.com/goreleaser/goreleaser/pkg/context"
)

// Pipe for Artifactory
type Pipe struct{}

// String returns the description of the pipe
func (Pipe) String() string {
	return "releasing to s3"
}

// Default sets the pipe defaults
func (Pipe) Default(ctx *context.Context) error {
	for i := range ctx.Config.S3 {
		s3 := &ctx.Config.S3[i]
		if s3.Bucket == "" {
			continue
		}
		if s3.Folder == "" {
			s3.Folder = "{{ .ProjectName }}/{{ .Tag }}"
		}
		if s3.Region == "" {
			s3.Region = "us-east-1"
		}
		if s3.ACL == "" {
			s3.ACL = "private"
		}
	}
	return nil
}

// Run the pipe
func (Pipe) Run(ctx *context.Context) error {
	if ctx.SkipPublish {
		return pipeline.ErrSkipPublishEnabled
	}
	var g = semerrgroup.New(ctx.Parallelism)
	for _, conf := range ctx.Config.S3 {
		conf := conf
		g.Go(func() error {
			return upload(ctx, conf)
		})
	}
	return g.Wait()
}

func upload(ctx *context.Context, conf config.S3) error {
	var awsConfig = aws.NewConfig()
	// TODO: add a test for this
	if conf.Endpoint != "" {
		awsConfig.Endpoint = aws.String(conf.Endpoint)
		awsConfig.S3ForcePathStyle = aws.Bool(true)
	}
	awsConfig.Credentials = credentials.NewChainCredentials([]credentials.Provider{
		&credentials.EnvProvider{},
		&credentials.SharedCredentialsProvider{
			Profile: conf.Profile,
		},
	})
	sess := session.Must(session.NewSession(awsConfig))
	svc := s3.New(sess, &aws.Config{
		Region: aws.String(conf.Region),
	})
	folder, err := tmpl.New(ctx).Apply(conf.Folder)
	if err != nil {
		return err
	}

	var g = semerrgroup.New(ctx.Parallelism)
	for _, artifact := range ctx.Artifacts.Filter(
		artifact.Or(
			artifact.ByType(artifact.UploadableArchive),
			artifact.ByType(artifact.UploadableBinary),
			artifact.ByType(artifact.Checksum),
			artifact.ByType(artifact.Signature),
			artifact.ByType(artifact.LinuxPackage),
		),
	).List() {
		artifact := artifact
		g.Go(func() error {
			f, err := os.Open(artifact.Path)
			if err != nil {
				return err
			}
			log.WithFields(log.Fields{
				"bucket":   conf.Bucket,
				"folder":   folder,
				"artifact": artifact.Name,
			}).Info("uploading")
			_, err = svc.PutObjectWithContext(ctx, &s3.PutObjectInput{
				Bucket: aws.String(conf.Bucket),
				Key:    aws.String(filepath.Join(folder, artifact.Name)),
				Body:   f,
				ACL:    aws.String(conf.ACL),
			})
			return err
		})
	}
	return g.Wait()
}
<commit_msg>feat: Add support for profile and assume an IAM role with MFA prompting for token code on stdin when uploading to s3<commit_after>// Package s3 provides a Pipe that push artifacts to s3/minio
package s3

import (
	"os"
	"path/filepath"

	"github.com/apex/log"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/goreleaser/goreleaser/internal/artifact"
	"github.com/goreleaser/goreleaser/internal/pipeline"
	"github.com/goreleaser/goreleaser/internal/semerrgroup"
	"github.com/goreleaser/goreleaser/internal/tmpl"
	"github.com/goreleaser/goreleaser/pkg/config"
	"github.com/goreleaser/goreleaser/pkg/context"
)

// Pipe for Artifactory
type Pipe struct{}

// String returns the description of the pipe
func (Pipe) String() string {
	return "releasing to s3"
}

// Default sets the pipe defaults
func (Pipe) Default(ctx *context.Context) error {
	for i := range ctx.Config.S3 {
		s3 := &ctx.Config.S3[i]
		if s3.Bucket == "" {
			continue
		}
		if s3.Folder == "" {
			s3.Folder = "{{ .ProjectName }}/{{ .Tag }}"
		}
		if s3.Region == "" {
			s3.Region = "us-east-1"
		}
		if s3.ACL == "" {
			s3.ACL = "private"
		}
	}
	return nil
}

// Run the pipe
func (Pipe) Run(ctx *context.Context) error {
	if ctx.SkipPublish {
		return pipeline.ErrSkipPublishEnabled
	}
	var g = semerrgroup.New(ctx.Parallelism)
	for _, conf := range ctx.Config.S3 {
		conf := conf
		g.Go(func() error {
			return upload(ctx, conf)
		})
	}
	return g.Wait()
}

func upload(ctx *context.Context, conf config.S3) error {
	var awsConfig = aws.NewConfig()
	// TODO: add a test for this
	if conf.Endpoint != "" {
		awsConfig.Endpoint = aws.String(conf.Endpoint)
		awsConfig.S3ForcePathStyle = aws.Bool(true)
	}
	awsConfig.Credentials = credentials.NewChainCredentials([]credentials.Provider{
		&credentials.EnvProvider{},
		&credentials.SharedCredentialsProvider{
			Profile: conf.Profile,
		},
	})

	_, err := awsConfig.Credentials.Get()
	var sess *session.Session
	if err == nil {
		sess = session.Must(session.NewSession(awsConfig))
	} else {
		// Specify profile and assume an IAM role with MFA prompting for token code on stdin
		sess = session.Must(session.NewSessionWithOptions(session.Options{
			AssumeRoleTokenProvider: stscreds.StdinTokenProvider,
			SharedConfigState:       session.SharedConfigEnable,
			Profile:                 conf.Profile,
		}))
	}

	svc := s3.New(sess, &aws.Config{
		Region: aws.String(conf.Region),
	})
	folder, err := tmpl.New(ctx).Apply(conf.Folder)
	if err != nil {
		return err
	}

	var g = semerrgroup.New(ctx.Parallelism)
	for _, artifact := range ctx.Artifacts.Filter(
		artifact.Or(
			artifact.ByType(artifact.UploadableArchive),
			artifact.ByType(artifact.UploadableBinary),
			artifact.ByType(artifact.Checksum),
			artifact.ByType(artifact.Signature),
			artifact.ByType(artifact.LinuxPackage),
		),
	).List() {
		artifact := artifact
		g.Go(func() error {
			f, err := os.Open(artifact.Path)
			if err != nil {
				return err
			}
			log.WithFields(log.Fields{
				"bucket":   conf.Bucket,
				"folder":   folder,
				"artifact": artifact.Name,
			}).Info("uploading")
			_, err = svc.PutObjectWithContext(ctx, &s3.PutObjectInput{
				Bucket: aws.String(conf.Bucket),
				Key:    aws.String(filepath.Join(folder, artifact.Name)),
				Body:   f,
				ACL:    aws.String(conf.ACL),
			})
			return err
		})
	}
	return g.Wait()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"github.com/hanwen/termite/cba"
	"github.com/hanwen/termite/termite"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"os/user"
	"runtime"
	"runtime/pprof"
	"syscall"
)

var _ = log.Printf

func handleStop(daemon *termite.Worker) {
	for {
		sig := <-signal.Incoming
		switch sig.(os.UnixSignal) {
		case syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGINT:
			log.Println("got signal: ", sig)
			req := termite.ShutdownRequest{Kill: true}
			rep := termite.ShutdownResponse{}
			daemon.Shutdown(&req, &rep)
		case syscall.SIGHUP:
			daemon.Report()
		}
	}
}

func main() {
	cachedir := flag.String("cachedir", "/var/cache/termite/worker-cache", "content cache")
	tmpdir := flag.String("tmpdir", "/var/tmp",
		"where to create FUSE mounts; should be on same partition as cachedir.")
	secretFile := flag.String("secret", "secret.txt", "file containing password.")
	port := flag.Int("port", 1232, "Start of port to try.")
	portRetry := flag.Int("port-retry", 10, "How many other ports to try.")
	coordinator := flag.String("coordinator", "", "Where to register the worker.")
	jobs := flag.Int("jobs", 1, "Max number of jobs to run.")
	reapcount := flag.Int("reap-count", 1, "Number of jobs per filesystem.")
	userFlag := flag.String("user", "nobody", "Run as this user.")
	memcache := flag.Int("filecache", 1024, "number of <128k files to cache in memory")
	logfile := flag.String("logfile", "", "Output log file to use.")
	stderrFile := flag.String("stderr", "", "File to write stderr output to.")
	paranoia := flag.Bool("paranoia", false, "Check attribute cache.")
	cpus := flag.Int("cpus", 1, "Number of CPUs to use.")
	heap := flag.Int("heap-size", 0, "Maximum heap size in MB.")
	cpuprofile := flag.String("profile", "", "File to write profile output to.")
	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if os.Geteuid() != 0 {
		log.Fatal("This program must run as root")
	}
	secret, err := ioutil.ReadFile(*secretFile)
	if err != nil {
		log.Fatal("ReadFile", err)
	}

	if *logfile != "" {
		f, err := os.OpenFile(*logfile, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0644)
		if err != nil {
			log.Fatal("Could not open log file.", err)
		}
		log.Println("Log output to", *logfile)
		log.SetOutput(f)
	} else {
		log.SetPrefix("W")
	}

	if *stderrFile != "" {
		stderr, err := os.OpenFile(*stderrFile, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0644)
		if err != nil {
			log.Fatal("Could not open stderr file", err)
		}
		err = syscall.Close(2)
		if err != nil {
			log.Fatalf("close stderr: %v", err)
		}
		_, err = syscall.Dup(stderr.Fd())
		if err != nil {
			log.Fatalf("dup: %v", err)
		}
		stderr.Close()
	}

	opts := termite.WorkerOptions{
		Secret:      secret,
		TempDir:     *tmpdir,
		Jobs:        *jobs,
		Paranoia:    *paranoia,
		ReapCount:   *reapcount,
		LogFileName: *logfile,
		StoreOptions: cba.StoreOptions{
			Dir:      *cachedir,
			MemCount: *memcache,
		},
		HeapLimit:   uint64(*heap) * (1 << 20),
		Coordinator: *coordinator,
		Port:        *port,
		PortRetry:   *portRetry,
	}
	if os.Geteuid() == 0 {
		nobody, err := user.Lookup(*userFlag)
		if err != nil {
			log.Fatalf("can't lookup %q: %v", *userFlag, err)
		}
		opts.User = nobody
	}

	daemon := termite.NewWorker(&opts)
	if *cpus > 0 {
		runtime.GOMAXPROCS(*cpus)
	}
	log.Printf("%s on %d CPUs", termite.Version(), runtime.GOMAXPROCS(0))
	go handleStop(daemon)
	daemon.RunWorkerServer()
}
<commit_msg>Write log files to unique files.<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/hanwen/termite/cba"
	"github.com/hanwen/termite/termite"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"os/user"
	"runtime"
	"runtime/pprof"
	"syscall"
)

var _ = log.Printf

func handleStop(daemon *termite.Worker) {
	for {
		sig := <-signal.Incoming
		switch sig.(os.UnixSignal) {
		case syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGINT:
			log.Println("got signal: ", sig)
			req := termite.ShutdownRequest{Kill: true}
			rep := termite.ShutdownResponse{}
			daemon.Shutdown(&req, &rep)
		case syscall.SIGHUP:
			daemon.Report()
		}
	}
}

func OpenUniqueLog(base string) *os.File {
	name := base
	i := 0
	for {
		fi, _ := os.Stat(name)
		if fi == nil {
			break
		}

		name = fmt.Sprintf("%s.%d", base, i)
		i++
	}

	f, err := os.OpenFile(name, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		log.Fatal("Could not open log file.", err)
	}
	return f
}

func main() {
	cachedir := flag.String("cachedir", "/var/cache/termite/worker-cache", "content cache")
	tmpdir := flag.String("tmpdir", "/var/tmp",
		"where to create FUSE mounts; should be on same partition as cachedir.")
	secretFile := flag.String("secret", "secret.txt", "file containing password.")
	port := flag.Int("port", 1232, "Start of port to try.")
	portRetry := flag.Int("port-retry", 10, "How many other ports to try.")
	coordinator := flag.String("coordinator", "", "Where to register the worker.")
	jobs := flag.Int("jobs", 1, "Max number of jobs to run.")
	reapcount := flag.Int("reap-count", 1, "Number of jobs per filesystem.")
	userFlag := flag.String("user", "nobody", "Run as this user.")
	memcache := flag.Int("filecache", 1024, "number of <128k files to cache in memory")
	logfile := flag.String("logfile", "", "Output log file to use.")
	stderrFile := flag.String("stderr", "", "File to write stderr output to.")
	paranoia := flag.Bool("paranoia", false, "Check attribute cache.")
	cpus := flag.Int("cpus", 1, "Number of CPUs to use.")
	heap := flag.Int("heap-size", 0, "Maximum heap size in MB.")
	cpuprofile := flag.String("profile", "", "File to write profile output to.")
	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if os.Geteuid() != 0 {
		log.Fatal("This program must run as root")
	}
	secret, err := ioutil.ReadFile(*secretFile)
	if err != nil {
		log.Fatal("ReadFile", err)
	}

	if *logfile != "" {
		f := OpenUniqueLog(*logfile)
		log.Println("Log output to", *logfile)
		log.SetOutput(f)
	} else {
		log.SetPrefix("W")
	}

	if *stderrFile != "" {
		f := OpenUniqueLog(*stderrFile)
		err = syscall.Close(2)
		if err != nil {
			log.Fatalf("close stderr: %v", err)
		}
		_, err = syscall.Dup(f.Fd())
		if err != nil {
			log.Fatalf("dup: %v", err)
		}
		f.Close()
	}

	opts := termite.WorkerOptions{
		Secret:      secret,
		TempDir:     *tmpdir,
		Jobs:        *jobs,
		Paranoia:    *paranoia,
		ReapCount:   *reapcount,
		LogFileName: *logfile,
		StoreOptions: cba.StoreOptions{
			Dir:      *cachedir,
			MemCount: *memcache,
		},
		HeapLimit:   uint64(*heap) * (1 << 20),
		Coordinator: *coordinator,
		Port:        *port,
		PortRetry:   *portRetry,
	}
	if os.Geteuid() == 0 {
		nobody, err := user.Lookup(*userFlag)
		if err != nil {
			log.Fatalf("can't lookup %q: %v", *userFlag, err)
		}
		opts.User = nobody
	}

	daemon := termite.NewWorker(&opts)
	if *cpus > 0 {
		runtime.GOMAXPROCS(*cpus)
	}
	log.Printf("%s on %d CPUs", termite.Version(), runtime.GOMAXPROCS(0))
	go handleStop(daemon)
	daemon.RunWorkerServer()
}
<|endoftext|>


<commit_before>package server

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

	"github.com/go-chi/chi"
	"github.com/patrickmn/go-cache"
	"github.com/zerowidth/gh-shorthand/internal/pkg/config"
)

const (
	resultTTL     = 10 * time.Minute
	errorTTL      = 10 * time.Second
	sweepInterval = 10 * time.Minute
)

// RPCHandler is a set of RPC http handlers
type RPCHandler struct {
	cfg   *config.Config
	cache *cache.Cache
}

// NewRPCHandler creates a new RPC server with the given config
func NewRPCHandler(cfg *config.Config) *RPCHandler {
	return &RPCHandler{
		cfg:   cfg,
		cache: cache.New(resultTTL, sweepInterval),
	}
}

// Mount routes the RPC handlers on a mux
func (rpc *RPCHandler) Mount(mux *chi.Mux) {
	mux.Get("/", rpc.testHandler)
}

func (rpc *RPCHandler) testHandler(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		http.Error(w, http.StatusText(400), 400)
		return
	}

	query := r.Form.Get("q")
	if len(query) == 0 {
		return
	}

	var req Request
	cr, ok := rpc.cache.Get(query)
	if ok {
		req = *cr.(*Request)
	} else {
		req = Request{}
		go rpc.makeRequest(&req, query)
		rpc.cache.Set(query, &req, resultTTL)
	}

	if req.Done() {
		if err := req.Error(); err != nil {
			w.WriteHeader(400)
			fmt.Fprintf(w, "request error: %s", err.Error())
			return
		}

		fmt.Fprintf(w, "request value: %s", req.Value())
		return
	}

	w.WriteHeader(204) // No Content (request is pending)
}

func (rpc *RPCHandler) makeRequest(req *Request, query string) {
	log.Println("RPC request: ", query)
	<-time.After(3 * time.Second)
	if query == "error" {
		log.Println("RPC request error: ", query)
		req.setError(fmt.Errorf("an error occurred in the rpc service"))
		return
	}
	log.Println("RPC result: ", query)
	req.setValue(fmt.Sprintf("RPC sees: %s", query))
}

// Request represents a pending or completed RPC request
type Request struct {
	done  bool
	value string
	error error
	m     sync.RWMutex
}

// Done checks if this RPC request is done yet
func (r *Request) Done() bool {
	r.m.RLock()
	defer r.m.RUnlock()
	return r.done
}

// Value retrieves the request's value, if applicable
func (r *Request) Value() string {
	r.m.RLock()
	defer r.m.RUnlock()
	return r.value
}

// Error retrieves the request's error, if applicable
func (r *Request) Error() error {
	r.m.RLock()
	defer r.m.RUnlock()
	return r.error
}

func (r *Request) setValue(v string) {
	r.m.Lock()
	r.value = v
	r.done = true
	r.m.Unlock()
}

func (r *Request) setError(e error) {
	r.m.Lock()
	r.error = e
	r.done = true
	r.m.Unlock()
}
<commit_msg>Handle pending and completed requests separately<commit_after>package server

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

	"github.com/go-chi/chi"
	"github.com/patrickmn/go-cache"
	"github.com/zerowidth/gh-shorthand/internal/pkg/config"
)

const (
	resultTTL     = 10 * time.Minute
	errorTTL      = 10 * time.Second
	sweepInterval = 10 * time.Minute
)

// RPCHandler is a set of RPC http handlers
type RPCHandler struct {
	cfg     *config.Config
	cache   *cache.Cache
	m       sync.Mutex
	pending map[string]struct{}
}

// NewRPCHandler creates a new RPC server with the given config
func NewRPCHandler(cfg *config.Config) *RPCHandler {
	handler := RPCHandler{
		cfg:     cfg,
		cache:   cache.New(resultTTL, sweepInterval),
		pending: make(map[string]struct{}),
	}
	return &handler
}

// Mount routes the RPC handlers on a mux
func (rpc *RPCHandler) Mount(mux *chi.Mux) {
	mux.Get("/", rpc.testHandler)
}

func (rpc *RPCHandler) testHandler(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		http.Error(w, http.StatusText(400), 400)
		return
	}

	query := r.Form.Get("q")
	if len(query) == 0 {
		return
	}

	// now that basic checks are done, lock the cache and pending map to see
	// if the request is already in flight.
	rpc.m.Lock()
	defer rpc.m.Unlock()
	var res Result

	if _, ok := rpc.pending[query]; ok {
		w.WriteHeader(204) // No Content (request is pending)
		return
	} else if cr, ok := rpc.cache.Get(query); ok {
		res = cr.(Result)
		if res.Error != nil {
			w.WriteHeader(400) // bad request
			fmt.Fprintf(w, "request error: %s", res.Error.Error())
			return
		}
		fmt.Fprintf(w, "request value: %s", res.Value)
		return
	}
	go rpc.makeRequest(query)
	w.WriteHeader(204) // No Content (request is pending)
}

func (rpc *RPCHandler) makeRequest(query string) {
	rpc.m.Lock()
	rpc.pending[query] = struct{}{}
	rpc.m.Unlock()

	log.Println("RPC request: ", query)
	<-time.After(2 * time.Second)
	res := Result{Query: query}
	var ttl time.Duration
	if query == "error" {
		log.Println("RPC request error: ", query)
		res.Error = fmt.Errorf("an error occurred in the rpc service")
		ttl = errorTTL
	} else {
		log.Println("RPC result: ", query)
		res.Value = fmt.Sprintf("RPC sees: %s", query)
		ttl = resultTTL
	}

	rpc.m.Lock()
	delete(rpc.pending, query)
	rpc.cache.Set(query, res, ttl)
	rpc.m.Unlock()
}

// Result is the result of an RPC call
type Result struct {
	Query string
	Value string
	Error error
}
<|endoftext|>


<commit_before>// Copyright 2019 Authors of Cilium
//
// 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"
	"net"
	"strconv"
	"strings"
	"time"

	"github.com/cilium/cilium/api/v1/models"
	. "github.com/cilium/cilium/api/v1/server/restapi/policy"
	"github.com/cilium/cilium/pkg/api"
	"github.com/cilium/cilium/pkg/endpoint"
	"github.com/cilium/cilium/pkg/endpointmanager"
	"github.com/cilium/cilium/pkg/fqdn"
	"github.com/cilium/cilium/pkg/fqdn/dnsproxy"
	"github.com/cilium/cilium/pkg/fqdn/matchpattern"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/node"
	"github.com/cilium/cilium/pkg/option"
	policyApi "github.com/cilium/cilium/pkg/policy/api"
	"github.com/cilium/cilium/pkg/proxy"
	"github.com/cilium/cilium/pkg/proxy/accesslog"
	"github.com/cilium/cilium/pkg/proxy/logger"
	"github.com/cilium/cilium/pkg/u8proto"
	"github.com/go-openapi/runtime/middleware"
	"github.com/go-openapi/strfmt"

	"github.com/miekg/dns"
)

// bootstrapFQDN initializes the toFQDNs related subsystems: DNSPoller,
// d.dnsRuleGen, and the DNS proxy.
// dnsRuleGen and DNSPoller will use the default resolver and, implicitly, the
// default DNS cache. The proxy binds to all interfaces, and uses the
// configured DNS proxy port (this may be 0 and so OS-assigned).
func (d *Daemon) bootstrapFQDN(restoredEndpoints *endpointRestoreState) (err error) {
	cfg := fqdn.Config{
		MinTTL:         option.Config.ToFQDNsMinTTL,
		Cache:          fqdn.DefaultDNSCache,
		LookupDNSNames: fqdn.DNSLookupDefaultResolver,
		AddGeneratedRules: func(generatedRules []*policyApi.Rule) error {
			// Insert the new rules into the policy repository. We need them to
			// replace the previous set. This requires the labels to match (including
			// the ToFQDN-UUID one).
			_, err := d.PolicyAdd(generatedRules, &AddOptions{Replace: true, Generated: true})
			return err
		},
		PollerResponseNotify: func(lookupTime time.Time, qname string, response *fqdn.DNSIPRecords) {
			// Do nothing if this option is off
			if !option.Config.ToFQDNsEnablePollerEvents {
				return
			}

			// FIXME: Not always true but we don't have the protocol information here
			protocol := accesslog.TransportProtocol(u8proto.ProtoIDs["udp"])

			record := logger.LogRecord{
				LogRecord: accesslog.LogRecord{
					Type:              accesslog.TypeResponse,
					ObservationPoint:  accesslog.Ingress,
					IPVersion:         accesslog.VersionIPv4,
					TransportProtocol: protocol,
					Timestamp:         time.Now().UTC().Format(time.RFC3339Nano),
					NodeAddressInfo: accesslog.NodeAddressInfo{
						IPv4: node.GetExternalIPv4().String(),
						IPv6: node.GetIPv6().String(),
					},
				},
			}

			logger.LogTags.Verdict(accesslog.VerdictForwarded, "DNSPoller")(&record)
			logger.LogTags.DNS(&accesslog.LogRecordDNS{
				Query:             qname,
				IPs:               response.IPs,
				TTL:               uint32(response.TTL),
				CNAMEs:            nil,
				ObservationSource: accesslog.DNSSourceAgentPoller,
			})(&record)
			record.Log()
		}}

	d.dnsRuleGen = fqdn.NewRuleGen(cfg)
	d.dnsPoller = fqdn.NewDNSPoller(cfg, d.dnsRuleGen)
	if option.Config.ToFQDNsEnablePoller {
		fqdn.StartDNSPoller(d.dnsPoller)
	}

	// Prefill the cache with DNS lookups from restored endpoints. This is needed
	// to maintain continuity of which IPs are allowed.
	// Note: This is TTL aware, and expired data will not be used (e.g. when
	// restoring after a long delay).
	for _, restoredEP := range restoredEndpoints.restored {
		// Upgrades from old ciliums have this nil
		if restoredEP.DNSHistory != nil {
			fqdn.DefaultDNSCache.UpdateFromCache(restoredEP.DNSHistory)
		}
	}

	// Once we stop returning errors from StartDNSProxy this should live in
	// StartProxySupport
	proxy.DefaultDNSProxy, err = dnsproxy.StartDNSProxy("", uint16(option.Config.ToFQDNsProxyPort),
		// LookupEPByIP
		func(endpointIP net.IP) (endpointID string, err error) {
			e := endpointmanager.LookupIPv4(endpointIP.String())
			if e == nil {
				return "", fmt.Errorf("Cannot find endpoint with IP %s", endpointIP.String())
			}

			return e.StringID(), nil
		},
		// NotifyOnDNSMsg handles DNS data in the daemon by emitting monitor
		// events, proxy metrics and storing DNS data in the DNS cache. This may
		// result in rule generation.
		// It will:
		// - Report a monitor error event and proxy metrics when the proxy sees an
		//   error, and when it can't process something in this function
		// - Report the verdict in a monitor event and emit proxy metrics
		// - Insert the DNS data into the cache when msg is a DNS response and we
		//   can lookup the endpoint related to it
		// srcAddr and dstAddr should match the packet reported on (i.e. the
		// endpoint is srcAddr for requests, and dstAddr for responses).
		func(lookupTime time.Time, srcAddr, dstAddr string, msg *dns.Msg, protocol string, allowed bool, proxyErr error) error {
			var protoID = u8proto.ProtoIDs[strings.ToLower(protocol)]

			var verdict accesslog.FlowVerdict
			var reason string
			switch {
			case proxyErr != nil:
				verdict = accesslog.VerdictError
				reason = "Error: " + proxyErr.Error()
			case allowed:
				verdict = accesslog.VerdictForwarded
				reason = "Allowed by policy"
			case !allowed:
				verdict = accesslog.VerdictDenied
				reason = "Denied by policy"
			}

			var epAddr string     // the address of the endpoint that originated the request
			var serverAddr string // the address of the DNS target
			var ingress = msg.Response
			var flowType accesslog.FlowType
			if ingress {
				flowType = accesslog.TypeResponse
				epAddr = dstAddr
				serverAddr = srcAddr
			} else {
				flowType = accesslog.TypeRequest
				epAddr = srcAddr
				serverAddr = dstAddr
			}

			var serverPort int
			_, serverPortStr, err := net.SplitHostPort(serverAddr)
			if err != nil {
				log.WithError(err).Error("cannot extract endpoint IP from DNS request")
			} else {
				if serverPort, err = strconv.Atoi(serverPortStr); err != nil {
					log.WithError(err).WithField(logfields.Port, serverPortStr).Error("cannot parse destination port")
				}
			}

			var ep *endpoint.Endpoint
			epIP, _, err := net.SplitHostPort(epAddr)
			if err != nil {
				log.WithError(err).Error("cannot extract endpoint IP from DNS request")
				ep.UpdateProxyStatistics("dns", uint16(serverPort), ingress, !ingress, accesslog.VerdictError)
				return err
			}
			ep = endpointmanager.LookupIPv4(epIP)
			if ep == nil {
				// This is a hard fail. We cannot proceed because record.Log requires a
				// non-nil ep, and we also don't want to insert this data into the
				// cache if we don't know that an endpoint asked for it (this is
				// asserted via ep != nil here and msg.Response && msg.Rcode ==
				// dns.RcodeSuccess below).
				err := fmt.Errorf("Cannot find matching endpoint for IPs %s or %s", srcAddr, dstAddr)
				log.WithError(err).Error("cannot find matching endpoint")
				return err
			}

			qname, responseIPs, TTL, CNAMEs, err := dnsproxy.ExtractMsgDetails(msg)
			if err != nil {
				// This error is ok because all these values are used for reporting, or filling in the cache.
				log.WithError(err).Error("cannot extract DNS message details")
			}

			ep.UpdateProxyStatistics("dns", uint16(serverPort), ingress, !ingress, verdict)
			record := logger.NewLogRecord(proxy.DefaultEndpointInfoRegistry, ep, flowType, ingress,
				func(lr *logger.LogRecord) { lr.LogRecord.TransportProtocol = accesslog.TransportProtocol(protoID) },
				logger.LogTags.Verdict(verdict, reason),
				logger.LogTags.Addressing(logger.AddressingInfo{
					SrcIPPort:   srcAddr,
					DstIPPort:   dstAddr,
					SrcIdentity: 0, // 0 more correctly finds src and dst EP data
				}),
				logger.LogTags.DNS(&accesslog.LogRecordDNS{
					Query:             qname,
					IPs:               responseIPs,
					TTL:               TTL,
					CNAMEs:            CNAMEs,
					ObservationSource: accesslog.DNSSourceProxy,
				}),
			)
			record.Log()

			if msg.Response && msg.Rcode == dns.RcodeSuccess {
				// This must happen before the ruleGen update below, to ensure that
				// this data is included in the serialized Endpoint object.
				// Note: We need to fixup minTTL to be consistent with how we insert it
				// elsewhere i.e. we don't want to lose the lower bound for DNS data
				// TTL if we reboot twice.
				log.WithField(logfields.EndpointID, ep.ID).Debug("Recording DNS lookup in endpoint specific cache")
				effectiveTTL := int(TTL)
				if effectiveTTL < option.Config.ToFQDNsMinTTL {
					effectiveTTL = option.Config.ToFQDNsMinTTL
				}
				ep.DNSHistory.Update(lookupTime, qname, responseIPs, effectiveTTL)
				log.Debug("Updating DNS name in cache from response to to query")
				err = d.dnsRuleGen.UpdateGenerateDNS(lookupTime, map[string]*fqdn.DNSIPRecords{
					qname: {
						IPs: responseIPs,
						TTL: int(effectiveTTL),
					}})
				if err != nil {
					log.WithError(err).Error("error updating internal DNS cache for rule generation")
				}
			}

			return nil
		})
	proxy.DefaultDNSProxy.SetRejectReply(option.Config.FQDNRejectResponse)
	return err // filled by StartDNSProxy
}

type getFqdnCache struct {
	daemon *Daemon
}

func NewGetFqdnCacheHandler(d *Daemon) GetFqdnCacheHandler {
	return &getFqdnCache{daemon: d}
}

func (h *getFqdnCache) Handle(params GetFqdnCacheParams) middleware.Responder {
	// endpoints we want data from
	endpoints := endpointmanager.GetEndpoints()

	CIDRStr := ""
	if params.Cidr != nil {
		CIDRStr = *params.Cidr
	}

	matchPatternStr := ""
	if params.Matchpattern != nil {
		matchPatternStr = *params.Matchpattern
	}

	lookups, err := extractDNSLookups(endpoints, CIDRStr, matchPatternStr)
	switch {
	case err != nil:
		return api.Error(GetFqdnCacheBadRequestCode, err)
	case len(lookups) == 0:
		return NewGetFqdnCacheIDNotFound()
	}

	return NewGetFqdnCacheIDOK().WithPayload(lookups)
}

type getFqdnCacheID struct {
	daemon *Daemon
}

func NewGetFqdnCacheIDHandler(d *Daemon) GetFqdnCacheIDHandler {
	return &getFqdnCacheID{daemon: d}
}

func (h *getFqdnCacheID) Handle(params GetFqdnCacheIDParams) middleware.Responder {
	var endpoints []*endpoint.Endpoint
	if params.ID != "" {
		ep, err := endpointmanager.Lookup(params.ID)
		switch {
		case err != nil:
			return api.Error(GetFqdnCacheIDBadRequestCode, err)
		case ep == nil:
			return api.Error(GetFqdnCacheIDNotFoundCode, fmt.Errorf("Cannot find endpoint %s", params.ID))
		default:
			endpoints = []*endpoint.Endpoint{ep}
		}
	}

	CIDRStr := ""
	if params.Cidr != nil {
		CIDRStr = *params.Cidr
	}

	matchPatternStr := ""
	if params.Matchpattern != nil {
		matchPatternStr = *params.Matchpattern
	}

	lookups, err := extractDNSLookups(endpoints, CIDRStr, matchPatternStr)
	switch {
	case err != nil:
		return api.Error(GetFqdnCacheBadRequestCode, err)
	case len(lookups) == 0:
		return NewGetFqdnCacheIDNotFound()
	}

	return NewGetFqdnCacheIDOK().WithPayload(lookups)
}

// extractDNSLookups returns API models.DNSLookup copies of DNS data in each
// endpoint's DNSHistory. These are filtered by CIDRStr and matchPatternStr if
// they are non-empty.
func extractDNSLookups(endpoints []*endpoint.Endpoint, CIDRStr, matchPatternStr string) (lookups []*models.DNSLookup, err error) {
	cidrMatcher := func(ip net.IP) bool { return true }
	if CIDRStr != "" {
		_, cidr, err := net.ParseCIDR(CIDRStr)
		if err != nil {
			return nil, err
		}
		cidrMatcher = func(ip net.IP) bool { return cidr.Contains(ip) }
	}

	nameMatcher := func(name string) bool { return true }
	if matchPatternStr != "" {
		matcher, err := matchpattern.Validate(matchPatternStr)
		if err != nil {
			return nil, err
		}
		nameMatcher = func(name string) bool { return matcher.MatchString(name) }
	}

	for _, ep := range endpoints {
		for _, lookup := range ep.DNSHistory.Dump() {
			if !nameMatcher(lookup.Name) {
				continue
			}

			// The API model needs strings
			IPStrings := make([]string, 0, len(lookup.IPs))

			// only proceed if any IP matches the cidr selector
			anIPMatches := false
			for _, ip := range lookup.IPs {
				anIPMatches = anIPMatches || cidrMatcher(ip)
				IPStrings = append(IPStrings, ip.String())
			}
			if !anIPMatches {
				continue
			}

			lookups = append(lookups, &models.DNSLookup{
				Fqdn:           lookup.Name,
				Ips:            IPStrings,
				LookupTime:     strfmt.DateTime(lookup.LookupTime),
				TTL:            int64(lookup.TTL),
				ExpirationTime: strfmt.DateTime(lookup.ExpirationTime),
				EndpointID:     int64(ep.ID),
			})
		}
	}

	return lookups, nil
}
<commit_msg>daemon: Sanitize fqdn/cache API matchpattern parameter<commit_after>// Copyright 2019 Authors of Cilium
//
// 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"
	"net"
	"strconv"
	"strings"
	"time"

	"github.com/cilium/cilium/api/v1/models"
	. "github.com/cilium/cilium/api/v1/server/restapi/policy"
	"github.com/cilium/cilium/pkg/api"
	"github.com/cilium/cilium/pkg/endpoint"
	"github.com/cilium/cilium/pkg/endpointmanager"
	"github.com/cilium/cilium/pkg/fqdn"
	"github.com/cilium/cilium/pkg/fqdn/dnsproxy"
	"github.com/cilium/cilium/pkg/fqdn/matchpattern"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/node"
	"github.com/cilium/cilium/pkg/option"
	policyApi "github.com/cilium/cilium/pkg/policy/api"
	"github.com/cilium/cilium/pkg/proxy"
	"github.com/cilium/cilium/pkg/proxy/accesslog"
	"github.com/cilium/cilium/pkg/proxy/logger"
	"github.com/cilium/cilium/pkg/u8proto"
	"github.com/go-openapi/runtime/middleware"
	"github.com/go-openapi/strfmt"

	"github.com/miekg/dns"
)

// bootstrapFQDN initializes the toFQDNs related subsystems: DNSPoller,
// d.dnsRuleGen, and the DNS proxy.
// dnsRuleGen and DNSPoller will use the default resolver and, implicitly, the
// default DNS cache. The proxy binds to all interfaces, and uses the
// configured DNS proxy port (this may be 0 and so OS-assigned).
func (d *Daemon) bootstrapFQDN(restoredEndpoints *endpointRestoreState) (err error) {
	cfg := fqdn.Config{
		MinTTL:         option.Config.ToFQDNsMinTTL,
		Cache:          fqdn.DefaultDNSCache,
		LookupDNSNames: fqdn.DNSLookupDefaultResolver,
		AddGeneratedRules: func(generatedRules []*policyApi.Rule) error {
			// Insert the new rules into the policy repository. We need them to
			// replace the previous set. This requires the labels to match (including
			// the ToFQDN-UUID one).
			_, err := d.PolicyAdd(generatedRules, &AddOptions{Replace: true, Generated: true})
			return err
		},
		PollerResponseNotify: func(lookupTime time.Time, qname string, response *fqdn.DNSIPRecords) {
			// Do nothing if this option is off
			if !option.Config.ToFQDNsEnablePollerEvents {
				return
			}

			// FIXME: Not always true but we don't have the protocol information here
			protocol := accesslog.TransportProtocol(u8proto.ProtoIDs["udp"])

			record := logger.LogRecord{
				LogRecord: accesslog.LogRecord{
					Type:              accesslog.TypeResponse,
					ObservationPoint:  accesslog.Ingress,
					IPVersion:         accesslog.VersionIPv4,
					TransportProtocol: protocol,
					Timestamp:         time.Now().UTC().Format(time.RFC3339Nano),
					NodeAddressInfo: accesslog.NodeAddressInfo{
						IPv4: node.GetExternalIPv4().String(),
						IPv6: node.GetIPv6().String(),
					},
				},
			}

			logger.LogTags.Verdict(accesslog.VerdictForwarded, "DNSPoller")(&record)
			logger.LogTags.DNS(&accesslog.LogRecordDNS{
				Query:             qname,
				IPs:               response.IPs,
				TTL:               uint32(response.TTL),
				CNAMEs:            nil,
				ObservationSource: accesslog.DNSSourceAgentPoller,
			})(&record)
			record.Log()
		}}

	d.dnsRuleGen = fqdn.NewRuleGen(cfg)
	d.dnsPoller = fqdn.NewDNSPoller(cfg, d.dnsRuleGen)
	if option.Config.ToFQDNsEnablePoller {
		fqdn.StartDNSPoller(d.dnsPoller)
	}

	// Prefill the cache with DNS lookups from restored endpoints. This is needed
	// to maintain continuity of which IPs are allowed.
	// Note: This is TTL aware, and expired data will not be used (e.g. when
	// restoring after a long delay).
	for _, restoredEP := range restoredEndpoints.restored {
		// Upgrades from old ciliums have this nil
		if restoredEP.DNSHistory != nil {
			fqdn.DefaultDNSCache.UpdateFromCache(restoredEP.DNSHistory)
		}
	}

	// Once we stop returning errors from StartDNSProxy this should live in
	// StartProxySupport
	proxy.DefaultDNSProxy, err = dnsproxy.StartDNSProxy("", uint16(option.Config.ToFQDNsProxyPort),
		// LookupEPByIP
		func(endpointIP net.IP) (endpointID string, err error) {
			e := endpointmanager.LookupIPv4(endpointIP.String())
			if e == nil {
				return "", fmt.Errorf("Cannot find endpoint with IP %s", endpointIP.String())
			}

			return e.StringID(), nil
		},
		// NotifyOnDNSMsg handles DNS data in the daemon by emitting monitor
		// events, proxy metrics and storing DNS data in the DNS cache. This may
		// result in rule generation.
		// It will:
		// - Report a monitor error event and proxy metrics when the proxy sees an
		//   error, and when it can't process something in this function
		// - Report the verdict in a monitor event and emit proxy metrics
		// - Insert the DNS data into the cache when msg is a DNS response and we
		//   can lookup the endpoint related to it
		// srcAddr and dstAddr should match the packet reported on (i.e. the
		// endpoint is srcAddr for requests, and dstAddr for responses).
		func(lookupTime time.Time, srcAddr, dstAddr string, msg *dns.Msg, protocol string, allowed bool, proxyErr error) error {
			var protoID = u8proto.ProtoIDs[strings.ToLower(protocol)]

			var verdict accesslog.FlowVerdict
			var reason string
			switch {
			case proxyErr != nil:
				verdict = accesslog.VerdictError
				reason = "Error: " + proxyErr.Error()
			case allowed:
				verdict = accesslog.VerdictForwarded
				reason = "Allowed by policy"
			case !allowed:
				verdict = accesslog.VerdictDenied
				reason = "Denied by policy"
			}

			var epAddr string     // the address of the endpoint that originated the request
			var serverAddr string // the address of the DNS target
			var ingress = msg.Response
			var flowType accesslog.FlowType
			if ingress {
				flowType = accesslog.TypeResponse
				epAddr = dstAddr
				serverAddr = srcAddr
			} else {
				flowType = accesslog.TypeRequest
				epAddr = srcAddr
				serverAddr = dstAddr
			}

			var serverPort int
			_, serverPortStr, err := net.SplitHostPort(serverAddr)
			if err != nil {
				log.WithError(err).Error("cannot extract endpoint IP from DNS request")
			} else {
				if serverPort, err = strconv.Atoi(serverPortStr); err != nil {
					log.WithError(err).WithField(logfields.Port, serverPortStr).Error("cannot parse destination port")
				}
			}

			var ep *endpoint.Endpoint
			epIP, _, err := net.SplitHostPort(epAddr)
			if err != nil {
				log.WithError(err).Error("cannot extract endpoint IP from DNS request")
				ep.UpdateProxyStatistics("dns", uint16(serverPort), ingress, !ingress, accesslog.VerdictError)
				return err
			}
			ep = endpointmanager.LookupIPv4(epIP)
			if ep == nil {
				// This is a hard fail. We cannot proceed because record.Log requires a
				// non-nil ep, and we also don't want to insert this data into the
				// cache if we don't know that an endpoint asked for it (this is
				// asserted via ep != nil here and msg.Response && msg.Rcode ==
				// dns.RcodeSuccess below).
				err := fmt.Errorf("Cannot find matching endpoint for IPs %s or %s", srcAddr, dstAddr)
				log.WithError(err).Error("cannot find matching endpoint")
				return err
			}

			qname, responseIPs, TTL, CNAMEs, err := dnsproxy.ExtractMsgDetails(msg)
			if err != nil {
				// This error is ok because all these values are used for reporting, or filling in the cache.
				log.WithError(err).Error("cannot extract DNS message details")
			}

			ep.UpdateProxyStatistics("dns", uint16(serverPort), ingress, !ingress, verdict)
			record := logger.NewLogRecord(proxy.DefaultEndpointInfoRegistry, ep, flowType, ingress,
				func(lr *logger.LogRecord) { lr.LogRecord.TransportProtocol = accesslog.TransportProtocol(protoID) },
				logger.LogTags.Verdict(verdict, reason),
				logger.LogTags.Addressing(logger.AddressingInfo{
					SrcIPPort:   srcAddr,
					DstIPPort:   dstAddr,
					SrcIdentity: 0, // 0 more correctly finds src and dst EP data
				}),
				logger.LogTags.DNS(&accesslog.LogRecordDNS{
					Query:             qname,
					IPs:               responseIPs,
					TTL:               TTL,
					CNAMEs:            CNAMEs,
					ObservationSource: accesslog.DNSSourceProxy,
				}),
			)
			record.Log()

			if msg.Response && msg.Rcode == dns.RcodeSuccess {
				// This must happen before the ruleGen update below, to ensure that
				// this data is included in the serialized Endpoint object.
				// Note: We need to fixup minTTL to be consistent with how we insert it
				// elsewhere i.e. we don't want to lose the lower bound for DNS data
				// TTL if we reboot twice.
				log.WithField(logfields.EndpointID, ep.ID).Debug("Recording DNS lookup in endpoint specific cache")
				effectiveTTL := int(TTL)
				if effectiveTTL < option.Config.ToFQDNsMinTTL {
					effectiveTTL = option.Config.ToFQDNsMinTTL
				}
				ep.DNSHistory.Update(lookupTime, qname, responseIPs, effectiveTTL)
				log.Debug("Updating DNS name in cache from response to to query")
				err = d.dnsRuleGen.UpdateGenerateDNS(lookupTime, map[string]*fqdn.DNSIPRecords{
					qname: {
						IPs: responseIPs,
						TTL: int(effectiveTTL),
					}})
				if err != nil {
					log.WithError(err).Error("error updating internal DNS cache for rule generation")
				}
			}

			return nil
		})
	proxy.DefaultDNSProxy.SetRejectReply(option.Config.FQDNRejectResponse)
	return err // filled by StartDNSProxy
}

type getFqdnCache struct {
	daemon *Daemon
}

func NewGetFqdnCacheHandler(d *Daemon) GetFqdnCacheHandler {
	return &getFqdnCache{daemon: d}
}

func (h *getFqdnCache) Handle(params GetFqdnCacheParams) middleware.Responder {
	// endpoints we want data from
	endpoints := endpointmanager.GetEndpoints()

	CIDRStr := ""
	if params.Cidr != nil {
		CIDRStr = *params.Cidr
	}

	matchPatternStr := ""
	if params.Matchpattern != nil {
		matchPatternStr = *params.Matchpattern
	}

	lookups, err := extractDNSLookups(endpoints, CIDRStr, matchPatternStr)
	switch {
	case err != nil:
		return api.Error(GetFqdnCacheBadRequestCode, err)
	case len(lookups) == 0:
		return NewGetFqdnCacheIDNotFound()
	}

	return NewGetFqdnCacheIDOK().WithPayload(lookups)
}

type getFqdnCacheID struct {
	daemon *Daemon
}

func NewGetFqdnCacheIDHandler(d *Daemon) GetFqdnCacheIDHandler {
	return &getFqdnCacheID{daemon: d}
}

func (h *getFqdnCacheID) Handle(params GetFqdnCacheIDParams) middleware.Responder {
	var endpoints []*endpoint.Endpoint
	if params.ID != "" {
		ep, err := endpointmanager.Lookup(params.ID)
		switch {
		case err != nil:
			return api.Error(GetFqdnCacheIDBadRequestCode, err)
		case ep == nil:
			return api.Error(GetFqdnCacheIDNotFoundCode, fmt.Errorf("Cannot find endpoint %s", params.ID))
		default:
			endpoints = []*endpoint.Endpoint{ep}
		}
	}

	CIDRStr := ""
	if params.Cidr != nil {
		CIDRStr = *params.Cidr
	}

	matchPatternStr := ""
	if params.Matchpattern != nil {
		matchPatternStr = *params.Matchpattern
	}

	lookups, err := extractDNSLookups(endpoints, CIDRStr, matchPatternStr)
	switch {
	case err != nil:
		return api.Error(GetFqdnCacheBadRequestCode, err)
	case len(lookups) == 0:
		return NewGetFqdnCacheIDNotFound()
	}

	return NewGetFqdnCacheIDOK().WithPayload(lookups)
}

// extractDNSLookups returns API models.DNSLookup copies of DNS data in each
// endpoint's DNSHistory. These are filtered by CIDRStr and matchPatternStr if
// they are non-empty.
func extractDNSLookups(endpoints []*endpoint.Endpoint, CIDRStr, matchPatternStr string) (lookups []*models.DNSLookup, err error) {
	cidrMatcher := func(ip net.IP) bool { return true }
	if CIDRStr != "" {
		_, cidr, err := net.ParseCIDR(CIDRStr)
		if err != nil {
			return nil, err
		}
		cidrMatcher = func(ip net.IP) bool { return cidr.Contains(ip) }
	}

	nameMatcher := func(name string) bool { return true }
	if matchPatternStr != "" {
		matcher, err := matchpattern.Validate(matchpattern.Sanitize(matchPatternStr))
		if err != nil {
			return nil, err
		}
		nameMatcher = func(name string) bool { return matcher.MatchString(name) }
	}

	for _, ep := range endpoints {
		for _, lookup := range ep.DNSHistory.Dump() {
			if !nameMatcher(lookup.Name) {
				continue
			}

			// The API model needs strings
			IPStrings := make([]string, 0, len(lookup.IPs))

			// only proceed if any IP matches the cidr selector
			anIPMatches := false
			for _, ip := range lookup.IPs {
				anIPMatches = anIPMatches || cidrMatcher(ip)
				IPStrings = append(IPStrings, ip.String())
			}
			if !anIPMatches {
				continue
			}

			lookups = append(lookups, &models.DNSLookup{
				Fqdn:           lookup.Name,
				Ips:            IPStrings,
				LookupTime:     strfmt.DateTime(lookup.LookupTime),
				TTL:            int64(lookup.TTL),
				ExpirationTime: strfmt.DateTime(lookup.ExpirationTime),
				EndpointID:     int64(ep.ID),
			})
		}
	}

	return lookups, nil
}
<|endoftext|>