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

package bam

import (
	"code.google.com/p/biogo.bam/bgzf"
	"code.google.com/p/biogo.bam/sam"

	"errors"
	"sort"
)

var baiMagic = [4]byte{'B', 'A', 'I', 0x1}

const (
	tileWidth     = 0x4000
	statsDummyBin = 0x924a
)

func makeOffset(vOff uint64) bgzf.Offset {
	return bgzf.Offset{
		File:  int64(vOff >> 16),
		Block: uint16(vOff),
	}
}

func isZero(o bgzf.Offset) bool {
	return o == bgzf.Offset{}
}

func vOffset(o bgzf.Offset) int64 {
	return o.File<<16 | int64(o.Block)
}

func isPlaced(r *sam.Record) bool {
	return r.Ref != nil && r.Pos != -1
}

type byBinNumber []Bin

func (b byBinNumber) Len() int           { return len(b) }
func (b byBinNumber) Less(i, j int) bool { return b[i].Bin < b[j].Bin }
func (b byBinNumber) Swap(i, j int)      { b[i], b[j] = b[j], b[i] }

type byBeginOffset []bgzf.Chunk

func (c byBeginOffset) Len() int           { return len(c) }
func (c byBeginOffset) Less(i, j int) bool { return vOffset(c[i].Begin) < vOffset(c[j].Begin) }
func (c byBeginOffset) Swap(i, j int)      { c[i], c[j] = c[j], c[i] }

type byVirtOffset []bgzf.Offset

func (o byVirtOffset) Len() int           { return len(o) }
func (o byVirtOffset) Less(i, j int) bool { return vOffset(o[i]) < vOffset(o[j]) }
func (o byVirtOffset) Swap(i, j int)      { o[i], o[j] = o[j], o[i] }

type Index struct {
	References []RefIndex
	Unmapped   *uint64
	isSorted   bool
	lastRecord int
}

type RefIndex struct {
	Bins      []Bin
	Stats     *IndexStats
	Intervals []bgzf.Offset
}

type Bin struct {
	Bin    uint32
	Chunks []bgzf.Chunk
}

type IndexStats struct {
	Chunk    bgzf.Chunk
	Mapped   uint64
	Unmapped uint64
}

func (i *Index) Sort() {
	if !i.isSorted {
		for _, ref := range i.References {
			sort.Sort(byBinNumber(ref.Bins))
			for _, bin := range ref.Bins {
				sort.Sort(byBeginOffset(bin.Chunks))
			}
			sort.Sort(byVirtOffset(ref.Intervals))
		}
		i.isSorted = true
	}
}

func (i *Index) Add(r *sam.Record, c bgzf.Chunk) error {
	if !validIndexPos(r.Start()) || !validIndexPos(r.End()) {
		return errors.New("bam: attempt to add record outside indexable range")
	}

	if i.Unmapped == nil {
		i.Unmapped = new(uint64)
	}
	if !isPlaced(r) {
		*i.Unmapped++
		return nil
	}

	rid := r.Ref.ID()
	if rid < len(i.References)-1 {
		return errors.New("bam: attempt to add record out of reference ID sort order")
	}
	if rid == len(i.References) {
		i.References = append(i.References, RefIndex{})
	} else {
		refs := make([]RefIndex, rid+1)
		copy(refs, i.References)
		i.References = refs
	}
	ref := &i.References[rid]

	// Record bin information.
	b := uint32(r.Bin())
	for i, bin := range ref.Bins {
		if bin.Bin == b {
			for j, chunk := range ref.Bins[i].Chunks {
				if vOffset(chunk.End) > vOffset(c.Begin) {
					ref.Bins[i].Chunks[j].End = c.End
					goto found
				}
			}
			ref.Bins[i].Chunks = append(ref.Bins[i].Chunks, c)
			goto found
		}
	}
	i.isSorted = false // TODO(kortschak) Consider making use of this more effectively for bin search.
	ref.Bins = append(ref.Bins, Bin{
		Bin:    b,
		Chunks: []bgzf.Chunk{c},
	})
found:

	// Record interval tile information.
	biv := r.Start() / tileWidth
	if r.Start() < i.lastRecord {
		return errors.New("bam: attempt to add record out of position sort order")
	}
	i.lastRecord = r.Start()
	eiv := r.End() / tileWidth
	if eiv == len(ref.Intervals) {
		if eiv > biv {
			panic("bam: unexpected alignment length")
		}
		ref.Intervals = append(ref.Intervals, c.Begin)
	} else if eiv > len(ref.Intervals) {
		intvs := make([]bgzf.Offset, eiv)
		if len(ref.Intervals) > biv {
			biv = len(ref.Intervals)
		}
		for iv, offset := range intvs[biv:eiv] {
			if !isZero(offset) {
				panic("bam: unexpected non-zero offset")
			}
			intvs[iv+biv] = c.Begin
		}
		copy(intvs, ref.Intervals)
		ref.Intervals = intvs
	}

	// Record index stats.
	if ref.Stats == nil {
		ref.Stats = &IndexStats{
			Chunk: c,
		}
	} else {
		ref.Stats.Chunk.End = c.End
	}
	if r.Flags&sam.Unmapped == 0 {
		ref.Stats.Mapped++
	} else {
		ref.Stats.Unmapped++
	}

	return nil
}

func (i *Index) Chunks(r *sam.Reference, beg, end int) []bgzf.Chunk {
	rid := r.ID()
	if rid < 0 || rid >= len(i.References) {
		return nil
	}
	i.Sort()
	ref := i.References[rid]

	iv := beg / tileWidth
	if iv >= len(ref.Intervals) {
		return nil
	}

	// Collect candidate chunks according to the scheme described in
	// the SAM spec under section 5 Indexing BAM.
	var chunks []bgzf.Chunk
	for _, bin := range reg2bins(beg, end) {
		b := uint32(bin)
		c := sort.Search(len(ref.Bins), func(i int) bool { return ref.Bins[i].Bin >= b })
		if c < len(ref.Bins) && ref.Bins[c].Bin == b {
			for _, chunk := range ref.Bins[c].Chunks {
				// Here we check all tiles starting from the left end of the
				// query region until we get a non-zero offset. The spec states
				// that we only need to check tiles that contain beg. That is
				// not correct since we may have no alignments at the left end
				// of the query region.
				for j, tile := range ref.Intervals[iv:] {
					if isZero(tile) {
						continue
					}
					tbeg := (j + iv) * tileWidth
					tend := tbeg + tileWidth
					// We allow adjacent alignment since samtools behaviour here
					// has always irritated me and it is cheap to discard these
					// later if they are not wanted.
					if tend >= beg && tbeg <= end && vOffset(chunk.End) > vOffset(tile) {
						chunks = append(chunks, chunk)
						break
					}
				}
			}
		}
	}

	// Sort and merge overlaps.
	if !sort.IsSorted(byBeginOffset(chunks)) {
		sort.Sort(byBeginOffset(chunks))
	}

	return adjacent(chunks)
}

type Strategy func([]bgzf.Chunk) []bgzf.Chunk

var (
	Identity Strategy = identity
	Adjacent Strategy = adjacent
	Squash   Strategy = squash
)

func CompressorStrategy(near int64) Strategy {
	return func(chunks []bgzf.Chunk) []bgzf.Chunk {
		if len(chunks) == 0 {
			return nil
		}
		for c := 1; c < len(chunks); c++ {
			leftChunk := chunks[c-1]
			rightChunk := &chunks[c]
			if leftChunk.End.File+near >= rightChunk.Begin.File {
				rightChunk.Begin = leftChunk.Begin
				if vOffset(leftChunk.End) > vOffset(rightChunk.End) {
					rightChunk.End = leftChunk.End
				}
				chunks = append(chunks[:c-1], chunks[c:]...)
				c--
			}
		}
		return chunks
	}
}

func identity(chunks []bgzf.Chunk) []bgzf.Chunk { return chunks }

func adjacent(chunks []bgzf.Chunk) []bgzf.Chunk {
	if len(chunks) == 0 {
		return nil
	}
	for c := 1; c < len(chunks); c++ {
		leftChunk := chunks[c-1]
		rightChunk := &chunks[c]
		leftEndOffset := vOffset(leftChunk.End)
		if leftEndOffset >= vOffset(rightChunk.Begin) {
			rightChunk.Begin = leftChunk.Begin
			if leftEndOffset > vOffset(rightChunk.End) {
				rightChunk.End = leftChunk.End
			}
			chunks = append(chunks[:c-1], chunks[c:]...)
			c--
		}
	}
	return chunks
}

func squash(chunks []bgzf.Chunk) []bgzf.Chunk {
	if len(chunks) == 0 {
		return nil
	}
	left := chunks[0].Begin
	right := chunks[0].End
	for _, c := range chunks[1:] {
		if vOffset(c.End) > vOffset(right) {
			right = c.End
		}
	}
	return []bgzf.Chunk{{Begin: left, End: right}}
}

func (i *Index) MergeChunks(s Strategy) {
	if s == nil {
		return
	}
	for _, ref := range i.References {
		for b, bin := range ref.Bins {
			if !sort.IsSorted(byBeginOffset(bin.Chunks)) {
				sort.Sort(byBeginOffset(bin.Chunks))
			}
			ref.Bins[b].Chunks = s(bin.Chunks)
			if !sort.IsSorted(byBeginOffset(bin.Chunks)) {
				sort.Sort(byBeginOffset(bin.Chunks))
			}
		}
	}
}
<commit_msg>Reorder indexing code<commit_after>// Copyright ©2014 The bíogo.bam Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package bam

import (
	"code.google.com/p/biogo.bam/bgzf"
	"code.google.com/p/biogo.bam/sam"

	"errors"
	"sort"
)

var baiMagic = [4]byte{'B', 'A', 'I', 0x1}

const (
	tileWidth     = 0x4000
	statsDummyBin = 0x924a
)

type Index struct {
	References []RefIndex
	Unmapped   *uint64
	isSorted   bool
	lastRecord int
}

type RefIndex struct {
	Bins      []Bin
	Stats     *IndexStats
	Intervals []bgzf.Offset
}

type Bin struct {
	Bin    uint32
	Chunks []bgzf.Chunk
}

type IndexStats struct {
	Chunk    bgzf.Chunk
	Mapped   uint64
	Unmapped uint64
}

func (i *Index) Add(r *sam.Record, c bgzf.Chunk) error {
	if !validIndexPos(r.Start()) || !validIndexPos(r.End()) {
		return errors.New("bam: attempt to add record outside indexable range")
	}

	if i.Unmapped == nil {
		i.Unmapped = new(uint64)
	}
	if !isPlaced(r) {
		*i.Unmapped++
		return nil
	}

	rid := r.Ref.ID()
	if rid < len(i.References)-1 {
		return errors.New("bam: attempt to add record out of reference ID sort order")
	}
	if rid == len(i.References) {
		i.References = append(i.References, RefIndex{})
	} else {
		refs := make([]RefIndex, rid+1)
		copy(refs, i.References)
		i.References = refs
	}
	ref := &i.References[rid]

	// Record bin information.
	b := uint32(r.Bin())
	for i, bin := range ref.Bins {
		if bin.Bin == b {
			for j, chunk := range ref.Bins[i].Chunks {
				if vOffset(chunk.End) > vOffset(c.Begin) {
					ref.Bins[i].Chunks[j].End = c.End
					goto found
				}
			}
			ref.Bins[i].Chunks = append(ref.Bins[i].Chunks, c)
			goto found
		}
	}
	i.isSorted = false // TODO(kortschak) Consider making use of this more effectively for bin search.
	ref.Bins = append(ref.Bins, Bin{
		Bin:    b,
		Chunks: []bgzf.Chunk{c},
	})
found:

	// Record interval tile information.
	biv := r.Start() / tileWidth
	if r.Start() < i.lastRecord {
		return errors.New("bam: attempt to add record out of position sort order")
	}
	i.lastRecord = r.Start()
	eiv := r.End() / tileWidth
	if eiv == len(ref.Intervals) {
		if eiv > biv {
			panic("bam: unexpected alignment length")
		}
		ref.Intervals = append(ref.Intervals, c.Begin)
	} else if eiv > len(ref.Intervals) {
		intvs := make([]bgzf.Offset, eiv)
		if len(ref.Intervals) > biv {
			biv = len(ref.Intervals)
		}
		for iv, offset := range intvs[biv:eiv] {
			if !isZero(offset) {
				panic("bam: unexpected non-zero offset")
			}
			intvs[iv+biv] = c.Begin
		}
		copy(intvs, ref.Intervals)
		ref.Intervals = intvs
	}

	// Record index stats.
	if ref.Stats == nil {
		ref.Stats = &IndexStats{
			Chunk: c,
		}
	} else {
		ref.Stats.Chunk.End = c.End
	}
	if r.Flags&sam.Unmapped == 0 {
		ref.Stats.Mapped++
	} else {
		ref.Stats.Unmapped++
	}

	return nil
}

func (i *Index) Chunks(r *sam.Reference, beg, end int) []bgzf.Chunk {
	rid := r.ID()
	if rid < 0 || rid >= len(i.References) {
		return nil
	}
	i.Sort()
	ref := i.References[rid]

	iv := beg / tileWidth
	if iv >= len(ref.Intervals) {
		return nil
	}

	// Collect candidate chunks according to the scheme described in
	// the SAM spec under section 5 Indexing BAM.
	var chunks []bgzf.Chunk
	for _, bin := range reg2bins(beg, end) {
		b := uint32(bin)
		c := sort.Search(len(ref.Bins), func(i int) bool { return ref.Bins[i].Bin >= b })
		if c < len(ref.Bins) && ref.Bins[c].Bin == b {
			for _, chunk := range ref.Bins[c].Chunks {
				// Here we check all tiles starting from the left end of the
				// query region until we get a non-zero offset. The spec states
				// that we only need to check tiles that contain beg. That is
				// not correct since we may have no alignments at the left end
				// of the query region.
				for j, tile := range ref.Intervals[iv:] {
					if isZero(tile) {
						continue
					}
					tbeg := (j + iv) * tileWidth
					tend := tbeg + tileWidth
					// We allow adjacent alignment since samtools behaviour here
					// has always irritated me and it is cheap to discard these
					// later if they are not wanted.
					if tend >= beg && tbeg <= end && vOffset(chunk.End) > vOffset(tile) {
						chunks = append(chunks, chunk)
						break
					}
				}
			}
		}
	}

	// Sort and merge overlaps.
	if !sort.IsSorted(byBeginOffset(chunks)) {
		sort.Sort(byBeginOffset(chunks))
	}

	return adjacent(chunks)
}

func (i *Index) Sort() {
	if !i.isSorted {
		for _, ref := range i.References {
			sort.Sort(byBinNumber(ref.Bins))
			for _, bin := range ref.Bins {
				sort.Sort(byBeginOffset(bin.Chunks))
			}
			sort.Sort(byVirtOffset(ref.Intervals))
		}
		i.isSorted = true
	}
}

func makeOffset(vOff uint64) bgzf.Offset {
	return bgzf.Offset{
		File:  int64(vOff >> 16),
		Block: uint16(vOff),
	}
}

func isZero(o bgzf.Offset) bool {
	return o == bgzf.Offset{}
}

func vOffset(o bgzf.Offset) int64 {
	return o.File<<16 | int64(o.Block)
}

func isPlaced(r *sam.Record) bool {
	return r.Ref != nil && r.Pos != -1
}

type byBinNumber []Bin

func (b byBinNumber) Len() int           { return len(b) }
func (b byBinNumber) Less(i, j int) bool { return b[i].Bin < b[j].Bin }
func (b byBinNumber) Swap(i, j int)      { b[i], b[j] = b[j], b[i] }

type byBeginOffset []bgzf.Chunk

func (c byBeginOffset) Len() int           { return len(c) }
func (c byBeginOffset) Less(i, j int) bool { return vOffset(c[i].Begin) < vOffset(c[j].Begin) }
func (c byBeginOffset) Swap(i, j int)      { c[i], c[j] = c[j], c[i] }

type byVirtOffset []bgzf.Offset

func (o byVirtOffset) Len() int           { return len(o) }
func (o byVirtOffset) Less(i, j int) bool { return vOffset(o[i]) < vOffset(o[j]) }
func (o byVirtOffset) Swap(i, j int)      { o[i], o[j] = o[j], o[i] }

type Strategy func([]bgzf.Chunk) []bgzf.Chunk

var (
	Identity Strategy = identity
	Adjacent Strategy = adjacent
	Squash   Strategy = squash
)

func CompressorStrategy(near int64) Strategy {
	return func(chunks []bgzf.Chunk) []bgzf.Chunk {
		if len(chunks) == 0 {
			return nil
		}
		for c := 1; c < len(chunks); c++ {
			leftChunk := chunks[c-1]
			rightChunk := &chunks[c]
			if leftChunk.End.File+near >= rightChunk.Begin.File {
				rightChunk.Begin = leftChunk.Begin
				if vOffset(leftChunk.End) > vOffset(rightChunk.End) {
					rightChunk.End = leftChunk.End
				}
				chunks = append(chunks[:c-1], chunks[c:]...)
				c--
			}
		}
		return chunks
	}
}

func identity(chunks []bgzf.Chunk) []bgzf.Chunk { return chunks }

func adjacent(chunks []bgzf.Chunk) []bgzf.Chunk {
	if len(chunks) == 0 {
		return nil
	}
	for c := 1; c < len(chunks); c++ {
		leftChunk := chunks[c-1]
		rightChunk := &chunks[c]
		leftEndOffset := vOffset(leftChunk.End)
		if leftEndOffset >= vOffset(rightChunk.Begin) {
			rightChunk.Begin = leftChunk.Begin
			if leftEndOffset > vOffset(rightChunk.End) {
				rightChunk.End = leftChunk.End
			}
			chunks = append(chunks[:c-1], chunks[c:]...)
			c--
		}
	}
	return chunks
}

func squash(chunks []bgzf.Chunk) []bgzf.Chunk {
	if len(chunks) == 0 {
		return nil
	}
	left := chunks[0].Begin
	right := chunks[0].End
	for _, c := range chunks[1:] {
		if vOffset(c.End) > vOffset(right) {
			right = c.End
		}
	}
	return []bgzf.Chunk{{Begin: left, End: right}}
}

func (i *Index) MergeChunks(s Strategy) {
	if s == nil {
		return
	}
	for _, ref := range i.References {
		for b, bin := range ref.Bins {
			if !sort.IsSorted(byBeginOffset(bin.Chunks)) {
				sort.Sort(byBeginOffset(bin.Chunks))
			}
			ref.Bins[b].Chunks = s(bin.Chunks)
			if !sort.IsSorted(byBeginOffset(bin.Chunks)) {
				sort.Sort(byBeginOffset(bin.Chunks))
			}
		}
	}
}
<|endoftext|>


<commit_before>package kinesis

import (
	"testing"
	"github.com/influxdb/telegraf/testutil"
	"github.com/stretchr/testify/require"
)

func TestFormatMetric(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration test in short mode")
	}

	k := &KinesisOutput{
		Format: "string",
	}

	p := testutil.MockBatchPoints().Points()[0]

	valid_string := "test1,tag1=value1 value=1 1257894000000000000"
	func_string, err := FormatMetric(k, p)

	if func_string != valid_string {
		t.Error("Expected ", valid_string)
	}
	require.NoError(t, err)

	k = &KinesisOutput{
		Format: "custom",
	}

	valid_custom := "test1,map[tag1:value1],test1,tag1=value1 value=1 1257894000000000000"
	func_custom, err := FormatMetric(k, p)

	if func_custom != valid_custom {
		t.Error("Expected ", valid_custom)
	}
	require.NoError(t, err)
}
<commit_msg>Go fmt kinesis output test file<commit_after>package kinesis

import (
	"github.com/influxdb/telegraf/testutil"
	"github.com/stretchr/testify/require"
	"testing"
)

func TestFormatMetric(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration test in short mode")
	}

	k := &KinesisOutput{
		Format: "string",
	}

	p := testutil.MockBatchPoints().Points()[0]

	valid_string := "test1,tag1=value1 value=1 1257894000000000000"
	func_string, err := FormatMetric(k, p)

	if func_string != valid_string {
		t.Error("Expected ", valid_string)
	}
	require.NoError(t, err)

	k = &KinesisOutput{
		Format: "custom",
	}

	valid_custom := "test1,map[tag1:value1],test1,tag1=value1 value=1 1257894000000000000"
	func_custom, err := FormatMetric(k, p)

	if func_custom != valid_custom {
		t.Error("Expected ", valid_custom)
	}
	require.NoError(t, err)
}
<|endoftext|>


<commit_before>package rpcserver_test

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"

	rs "github.com/tendermint/tendermint/rpc/lib/server"
	types "github.com/tendermint/tendermint/rpc/lib/types"
	"github.com/tendermint/tmlibs/log"
)

// Ensure that nefarious/unintended inputs to `params`
// do not crash our RPC handlers.
// See Issue https://github.com/tendermint/tendermint/issues/708.
func TestRPCParams(t *testing.T) {
	funcMap := map[string]*rs.RPCFunc{
		"c": rs.NewRPCFunc(func(s string, i int) (string, error) { return "foo", nil }, "s,i"),
	}
	mux := http.NewServeMux()
	buf := new(bytes.Buffer)
	logger := log.NewTMLogger(buf)
	rs.RegisterRPCFuncs(mux, funcMap, logger)

	tests := []struct {
		payload      string
		wantErr      string
		notification bool
	}{
		{`{"jsonrpc": "2.0"}`, "", true}, // The server SHOULD NOT respond to a notification according to JSONRPC Section 4.1.
		{`{"jsonrpc": "2.0", "id": "0"}`, "Method not found", false},
		{`{"jsonrpc": "2.0", "method": "y", "id": "0"}`, "Method not found", false},
		{`{"jsonrpc": "2.0", "method": "c", "id": "0", "params": null}`, "", false},
		{`{"method": "c", "id": "0", "params": {}}`, "", false},
		{`{"method": "c", "id": "0", "params": a}`, "invalid character", false},
		{`{"method": "c", "id": "0", "params": ["a", 10]}`, "", false},
		{`{"method": "c", "id": "0", "params": ["a"]}`, "got 1", false},
		{`{"method": "c", "id": "0", "params": ["a", "b"]}`, "of type int", false},
		{`{"method": "c", "id": "0", "params": [1, 1]}`, "of type string", false},
	}

	statusOK := func(code int) bool { return code >= 200 && code <= 299 }

	for i, tt := range tests {
		req, _ := http.NewRequest("POST", "http://localhost/", strings.NewReader(tt.payload))
		rec := httptest.NewRecorder()
		mux.ServeHTTP(rec, req)
		res := rec.Result()
		// Always expecting back a JSONRPCResponse
		assert.True(t, statusOK(res.StatusCode), "#%d: should always return 2XX", i)
		blob, err := ioutil.ReadAll(res.Body)
		if err != nil {
			t.Errorf("#%d: err reading body: %v", i, err)
			continue
		}

		if tt.notification {
			assert.Equal(t, len(blob), 0, "#%d: a notification SHOULD NOT be responded to by the server", i)
			continue
		}

		recv := new(types.RPCResponse)
		assert.Nil(t, json.Unmarshal(blob, recv), "#%d: expecting successful parsing of an RPCResponse:\nblob: %s", i, blob)
		assert.NotEqual(t, recv, new(types.RPCResponse), "#%d: not expecting a blank RPCResponse", i)

		if tt.wantErr == "" {
			assert.Nil(t, recv.Error, "#%d: not expecting an error", i)
		} else {
			assert.False(t, statusOK(recv.Error.Code), "#%d: not expecting a 2XX success code", i)
			// The wanted error is either in the message or the data
			assert.Contains(t, recv.Error.Message+recv.Error.Data, tt.wantErr, "#%d: expected substring", i)
		}
	}
}
<commit_msg>rpc/lib/server: separate out Notifications test<commit_after>package rpcserver_test

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

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

	rs "github.com/tendermint/tendermint/rpc/lib/server"
	types "github.com/tendermint/tendermint/rpc/lib/types"
	"github.com/tendermint/tmlibs/log"
)

func testMux() *http.ServeMux {
	funcMap := map[string]*rs.RPCFunc{
		"c": rs.NewRPCFunc(func(s string, i int) (string, error) { return "foo", nil }, "s,i"),
	}
	mux := http.NewServeMux()
	buf := new(bytes.Buffer)
	logger := log.NewTMLogger(buf)
	rs.RegisterRPCFuncs(mux, funcMap, logger)

	return mux
}

func statusOK(code int) bool { return code >= 200 && code <= 299 }

// Ensure that nefarious/unintended inputs to `params`
// do not crash our RPC handlers.
// See Issue https://github.com/tendermint/tendermint/issues/708.
func TestRPCParams(t *testing.T) {
	mux := testMux()
	tests := []struct {
		payload string
		wantErr string
	}{
		{`{"jsonrpc": "2.0", "id": "0"}`, "Method not found"},
		{`{"jsonrpc": "2.0", "method": "y", "id": "0"}`, "Method not found"},
		{`{"jsonrpc": "2.0", "method": "c", "id": "0", "params": null}`, ""},
		{`{"method": "c", "id": "0", "params": {}}`, ""},
		{`{"method": "c", "id": "0", "params": a}`, "invalid character"},
		{`{"method": "c", "id": "0", "params": ["a", 10]}`, ""},
		{`{"method": "c", "id": "0", "params": ["a"]}`, "got 1"},
		{`{"method": "c", "id": "0", "params": ["a", "b"]}`, "of type int"},
		{`{"method": "c", "id": "0", "params": [1, 1]}`, "of type string"},
	}

	for i, tt := range tests {
		req, _ := http.NewRequest("POST", "http://localhost/", strings.NewReader(tt.payload))
		rec := httptest.NewRecorder()
		mux.ServeHTTP(rec, req)
		res := rec.Result()
		// Always expecting back a JSONRPCResponse
		assert.True(t, statusOK(res.StatusCode), "#%d: should always return 2XX", i)
		blob, err := ioutil.ReadAll(res.Body)
		if err != nil {
			t.Errorf("#%d: err reading body: %v", i, err)
			continue
		}

		recv := new(types.RPCResponse)
		assert.Nil(t, json.Unmarshal(blob, recv), "#%d: expecting successful parsing of an RPCResponse:\nblob: %s", i, blob)
		assert.NotEqual(t, recv, new(types.RPCResponse), "#%d: not expecting a blank RPCResponse", i)

		if tt.wantErr == "" {
			assert.Nil(t, recv.Error, "#%d: not expecting an error", i)
		} else {
			assert.False(t, statusOK(recv.Error.Code), "#%d: not expecting a 2XX success code", i)
			// The wanted error is either in the message or the data
			assert.Contains(t, recv.Error.Message+recv.Error.Data, tt.wantErr, "#%d: expected substring", i)
		}
	}
}

func TestRPCNotification(t *testing.T) {
	mux := testMux()
	body := strings.NewReader(`{"jsonrpc": "2.0"}`)
	req, _ := http.NewRequest("POST", "http://localhost/", body)
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	res := rec.Result()

	// Always expecting back a JSONRPCResponse
	require.True(t, statusOK(res.StatusCode), "should always return 2XX")
	blob, err := ioutil.ReadAll(res.Body)
	require.Nil(t, err, "reading from the body should not give back an error")
	require.Equal(t, len(blob), 0, "a notification SHOULD NOT be responded to by the server")
}
<|endoftext|>


<commit_before>package client

import (
	acm "github.com/eris-ltd/eris-db/account"
	core_types "github.com/eris-ltd/eris-db/core/types"
	rpc_types "github.com/eris-ltd/eris-db/rpc/tendermint/core/types"
	"github.com/eris-ltd/eris-db/txs"
	rpcclient "github.com/tendermint/go-rpc/client"
	"github.com/tendermint/go-wire"
)

func Status(client rpcclient.Client) (*rpc_types.ResultStatus, error) {
	var res rpc_types.ErisDBResult //ResultStatus)
	var err error
	switch cli := client.(type) {
	case *rpcclient.ClientJSONRPC:
		_, err = cli.Call("status", []interface{}{}, &res)
	case *rpcclient.ClientURI:
		_, err = cli.Call("status", map[string]interface{}{}, &res)
	}
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultStatus), nil
}

func GenPrivAccount(client rpcclient.Client) (*acm.PrivAccount, error) {
	var res rpc_types.ErisDBResult
	var err error
	switch cli := client.(type) {
	case *rpcclient.ClientJSONRPC:
		_, err = cli.Call("unsafe/gen_priv_account", []interface{}{}, &res)
	case *rpcclient.ClientURI:
		_, err = cli.Call("unsafe/gen_priv_account", map[string]interface{}{}, &res)
	}
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultGenPrivAccount).PrivAccount, nil
}

func GetAccount(client rpcclient.Client, addr []byte) (*acm.Account, error) {
	var res rpc_types.ErisDBResult
	var err error
	switch cli := client.(type) {
	case *rpcclient.ClientJSONRPC:
		_, err = cli.Call("get_account", []interface{}{addr}, &res)
	case *rpcclient.ClientURI:
		_, err = cli.Call("get_account", map[string]interface{}{"address": addr},
			&res)
	}
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultGetAccount).Account, nil
}

func SignTx(client rpcclient.Client, tx txs.Tx,
	privAccs []*acm.PrivAccount) (txs.Tx, error) {
	wrapTx := struct {
		txs.Tx `json:"unwrap"`
	}{tx}
	var res rpc_types.ErisDBResult
	var err error
	switch cli := client.(type) {
	case *rpcclient.ClientJSONRPC:
		_, err = cli.Call("unsafe/sign_tx", []interface{}{wrapTx, privAccs}, &res)
	case *rpcclient.ClientURI:
		_, err = cli.Call("unsafe/sign_tx", map[string]interface{}{
			"tx":           wrapTx,
			"privAccounts": privAccs,
		}, &res)
	}
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultSignTx).Tx, nil
}

func BroadcastTx(client rpcclient.Client,
	tx txs.Tx) (txs.Receipt, error) {
	wrapTx := struct {
		txs.Tx `json:"unwrap"`
	}{tx}
	var res rpc_types.ErisDBResult
	var err error
	switch cli := client.(type) {
	case *rpcclient.ClientJSONRPC:
		_, err = cli.Call("broadcast_tx", []interface{}{wrapTx}, &res)
	case *rpcclient.ClientURI:
		_, err = cli.Call("broadcast_tx", map[string]interface{}{"tx": wrapTx},
			&res)
	}
	if err != nil {
		return txs.Receipt{}, err
	}
	receiptBytes := res.(*rpc_types.ResultBroadcastTx).Data
	receipt := txs.Receipt{}
	err = wire.ReadBinaryBytes(receiptBytes, &receipt)
	return receipt, err

}

func DumpStorage(client rpcclient.Client,
	addr []byte) (*rpc_types.ResultDumpStorage, error) {
	var res rpc_types.ErisDBResult
	var err error
	switch cli := client.(type) {
	case *rpcclient.ClientJSONRPC:
		_, err = cli.Call("dump_storage", []interface{}{addr}, &res)
	case *rpcclient.ClientURI:
		_, err = cli.Call("dump_storage", map[string]interface{}{"address": addr},
			&res)
	}
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultDumpStorage), err
}

func GetStorage(client rpcclient.Client, addr, key []byte) ([]byte, error) {
	var res rpc_types.ErisDBResult
	var err error
	switch cli := client.(type) {
	case *rpcclient.ClientJSONRPC:
		_, err = cli.Call("get_storage", []interface{}{addr, key}, &res)
	case *rpcclient.ClientURI:
		_, err = cli.Call("get_storage", map[string]interface{}{
			"address": addr,
			"key":     key,
		}, &res)
	}
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultGetStorage).Value, nil
}

func CallCode(client rpcclient.Client,
	fromAddress, code, data []byte) (*rpc_types.ResultCall, error) {
	var res rpc_types.ErisDBResult
	var err error
	switch cli := client.(type) {
	case *rpcclient.ClientJSONRPC:
		_, err = cli.Call("call_code", []interface{}{fromAddress, code, data}, &res)
	case *rpcclient.ClientURI:
		_, err = cli.Call("call_code", map[string]interface{}{
			"fromAddress": fromAddress,
			"code":        code,
			"data":        data,
		}, &res)

	}
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultCall), err
}

func Call(client rpcclient.Client, fromAddress, toAddress,
	data []byte) (*rpc_types.ResultCall, error) {
	var res rpc_types.ErisDBResult
	var err error
	switch cli := client.(type) {
	case *rpcclient.ClientJSONRPC:
		_, err = cli.Call("call", []interface{}{fromAddress, toAddress, data}, &res)
	case *rpcclient.ClientURI:
		_, err = cli.Call("call", map[string]interface{}{
			"fromAddress": fromAddress,
			"toAddress":   toAddress,
			"data":        data,
		}, &res)

	}
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultCall), err
}

func GetName(client rpcclient.Client, name string) (*core_types.NameRegEntry, error) {
	var res rpc_types.ErisDBResult
	var err error
	switch cli := client.(type) {
	case *rpcclient.ClientJSONRPC:
		_, err = cli.Call("get_name", []interface{}{name}, &res)
	case *rpcclient.ClientURI:
		_, err = cli.Call("get_name", map[string]interface{}{"name": name}, &res)
	}
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultGetName).Entry, nil
}
<commit_msg>Client kept repeating itself it didso it did it did<commit_after>package client

import (
	"fmt"
	acm "github.com/eris-ltd/eris-db/account"
	core_types "github.com/eris-ltd/eris-db/core/types"
	rpc_types "github.com/eris-ltd/eris-db/rpc/tendermint/core/types"
	"github.com/eris-ltd/eris-db/txs"
	rpcclient "github.com/tendermint/go-rpc/client"
	"github.com/tendermint/go-wire"
)

func Status(client rpcclient.Client) (*rpc_types.ResultStatus, error) {
	res, err := performCall(client, "status")
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultStatus), nil
}

func GenPrivAccount(client rpcclient.Client) (*acm.PrivAccount, error) {
	res, err := performCall(client, "unsafe/gen_priv_account")
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultGenPrivAccount).PrivAccount, nil
}

func GetAccount(client rpcclient.Client, address []byte) (*acm.Account, error) {
	res, err := performCall(client, "get_account",
		"address", address)
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultGetAccount).Account, nil
}

func SignTx(client rpcclient.Client, tx txs.Tx,
	privAccounts []*acm.PrivAccount) (txs.Tx, error) {
	res, err := performCall(client, "unsafe/sign_tx",
		"tx", wrappedTx{tx},
		"privAccounts", privAccounts)
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultSignTx).Tx, nil
}

func BroadcastTx(client rpcclient.Client,
	tx txs.Tx) (txs.Receipt, error) {
	res, err := performCall(client, "broadcast_tx",
		"tx", wrappedTx{tx})
	if err != nil {
		return txs.Receipt{}, err
	}
	receiptBytes := res.(*rpc_types.ResultBroadcastTx).Data
	receipt := txs.Receipt{}
	err = wire.ReadBinaryBytes(receiptBytes, &receipt)
	return receipt, err

}

func DumpStorage(client rpcclient.Client,
	address []byte) (*rpc_types.ResultDumpStorage, error) {
	res, err := performCall(client, "dump_storage",
		"address", address)
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultDumpStorage), err
}

func GetStorage(client rpcclient.Client, address, key []byte) ([]byte, error) {
	res, err := performCall(client, "get_storage",
		"address", address,
		"key", key)
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultGetStorage).Value, nil
}

func CallCode(client rpcclient.Client, fromAddress, code,
	data []byte) (*rpc_types.ResultCall, error) {
	res, err := performCall(client, "call_code",
		"fromAddress", fromAddress,
		"code", code,
		"data", data)
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultCall), err
}

func Call(client rpcclient.Client, fromAddress, toAddress,
	data []byte) (*rpc_types.ResultCall, error) {
	res, err := performCall(client, "call",
		"fromAddress", fromAddress,
		"toAddress", toAddress,
		"data", data)
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultCall), err
}

func GetName(client rpcclient.Client, name string) (*core_types.NameRegEntry, error) {
	res, err := performCall(client, "get_name",
		"name", name)
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultGetName).Entry, nil
}

func BlockchainInfo(client rpcclient.Client, minHeight,
	maxHeight int) (*rpc_types.ResultBlockchainInfo, error) {
	res, err := performCall(client, "blockchain",
		"minHeight", minHeight,
		"maxHeight", maxHeight)
	if err != nil {
		return nil, err
	}
	return res.(*rpc_types.ResultBlockchainInfo), err
}

func performCall(client rpcclient.Client, method string,
	paramKeyVals ...interface{}) (res rpc_types.ErisDBResult, err error) {
	paramsMap, paramsSlice := mapAndValues(paramKeyVals...)
	switch cli := client.(type) {
	case *rpcclient.ClientJSONRPC:
		_, err = cli.Call(method, paramsSlice, &res)
	case *rpcclient.ClientURI:
		_, err = cli.Call(method, paramsMap, &res)
	}
	return res, err

}
func mapAndValues(orderedKeyVals ...interface{}) (map[string]interface{}, []interface{}) {
	if len(orderedKeyVals)%2 != 0 {
		panic(fmt.Sprintf("mapAndValues requires a even length list of keys and "+
			"values but got: %v (length %v)",
			orderedKeyVals, len(orderedKeyVals)))
	}
	paramsMap := make(map[string]interface{})
	paramsSlice := make([]interface{}, len(orderedKeyVals)/2)
	for i := 0; i < len(orderedKeyVals); i += 2 {
		key, ok := orderedKeyVals[i].(string)
		if !ok {
			panic("mapAndValues requires every even element of orderedKeyVals to be a string key")
		}
		val := orderedKeyVals[i+1]
		paramsMap[key] = val
		paramsSlice = append(paramsSlice, val)
	}
	return paramsMap, paramsSlice
}

type wrappedTx struct {
	txs.Tx `json:"unwrap"`
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/binary"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"math/rand"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/phil-mansfield/gotetra/los"
	"github.com/phil-mansfield/gotetra/los/geom"
	"github.com/phil-mansfield/gotetra/los/analyze"
	"github.com/phil-mansfield/gotetra/render/io"
	"github.com/phil-mansfield/gotetra/render/halo"
)

type Params struct {
	// HaloProfiles params
	RBins, Spokes, Rings int
	MaxMult, MinMult float64

	// Splashback params
	Order, Window, Levels int
	Cutoff float64
}

func main() {
	// Parse.
	p := parseCmd()
	ids, snaps, err := parseStdin()
	if err != nil { log.Fatal(err.Error()) }

	if len(ids) == 0 { return }

	// Compute coefficients.
	coeffs := make([][]float64, len(ids))
	snapBins, idxBins := binBySnap(snaps, ids)
	buf := make([]analyze.RingBuffer, p.Rings)
	for i := range buf { buf[i].Init(p.Spokes, p.RBins) }

	var losBuf *los.Buffers
	for snap, snapIDs := range snapBins {
		log.Println(snap)
		idxs := idxBins[snap]

		if snap == -1 { continue }

		// Bin halos
		hds, files, err := readHeaders(snap)
		if err != nil { err.Error() }
		if losBuf == nil { losBuf = los.NewBuffers(files[0], &hds[0]) }
		halos, err := createHalos(snap, &hds[0], snapIDs, p)
		if err != nil { log.Fatal(err.Error()) }
		intrBins := binIntersections(hds, halos)

		// Add densities. Done header by header to limit I/O time.
		for i := range hds {
			if len(intrBins[i]) == 0 { continue }
			hdContainer := []io.SheetHeader{hds[i]}
			fileContainer := []string{files[i]}
			los.LoadPtrDensities(
				intrBins[i], hdContainer, fileContainer, losBuf,
			)
		}

		// Calculate Penna coefficients.
		for i := range halos {
			coeffs[idxs[i]] = calcCoeffs(&halos[i], buf, p)
		}

	}
	printCoeffs(ids, snaps, coeffs)
}

func parseCmd() *Params {
	// Parse command line.
	p := &Params{}
	flag.IntVar(&p.RBins, "RBins", 256,
		"Number of radial bins used per LoS.")
	flag.IntVar(&p.Spokes, "Spokes", 1024,
		"Number of LoS's used per ring.")
	flag.IntVar(&p.Rings, "Rings", 10,
		"Number of rings used per halo. 3, 4, 6, and 10 rings are\n" +
			"guaranteed to be uniformly spaced.")
	flag.Float64Var(&p.MaxMult, "MaxMult", 3,
		"Ending radius of LoSs as a multiple of R_200m.")
	flag.Float64Var(&p.MinMult, "MinMult", 0.5,
		"Starting radius of LoSs as a multiple of R_200m.")
	flag.IntVar(&p.Order, "Order", 5,
		"Order of the shell fitting function.")
	flag.IntVar(&p.Window, "Window", 121,
		"Number of bins within smoothign window. Must be odd.")
	flag.IntVar(&p.Levels, "Levels", 4,
		"The number of recurve max-finding levels used by the 2D edge finder.")
	flag.Float64Var(&p.Cutoff, "Cutoff", 0.0,
		"The shallowest slope that can be considered a splashback point.")
	flag.Parse()
	return p
}

func parseStdin() (ids, snaps []int, err error) {
	ids, snaps = []int{}, []int{}
	lines, err := stdinLines()
	if err != nil { return nil, nil, err }
	for i, line := range lines {
		rawTokens := strings.Split(line, " ")
		tokens := make([]string, 0, len(rawTokens))
		for _, tok := range rawTokens {
			if len(tok) != 0 { tokens = append(tokens, tok) }
		}

		var (
			id, snap int
			err error
		)
		switch len(tokens) {
		case 0:
			continue
		case 2:
			id, err = strconv.Atoi(tokens[0])
			if err != nil {
				return nil, nil, fmt.Errorf(
					"One line %d of stdin, %s does not parse as an int.",
					i + 1, tokens[0],
				)
			}
			snap, err = strconv.Atoi(tokens[1])
			if err != nil {
				return nil, nil, fmt.Errorf(
					"One line %d of stdin, %s does not parse as an int.",
					i + 1, tokens[1],
				)
			}
		case 1:
			if tokens[0] == "" { continue }
			fallthrough
		default:
			return nil, nil, fmt.Errorf(
				"Line %d of stdin has %d tokens, but 2 are required.",
				i + 1, len(tokens),
			)
		}

		ids = append(ids, id)
		snaps = append(snaps, snap)
	}

	return ids, snaps, nil
}

func stdinLines() ([]string, error) {
		bs, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		return nil, fmt.Errorf(
			"Error reading stdin: %s.", err.Error(),
		)
	}

	text := string(bs)
	return strings.Split(text, "\n"), nil
}

func readHeaders(snap int) ([]io.SheetHeader, []string, error) {
	memoDir := os.Getenv("GTET_MEMO_DIR")
	if memoDir == "" {
		// You don't want to memoize? Fine. Deal with the consequences.
		return readHeadersFromSheet(snap)
	}
	if _, err := os.Stat(memoDir); err != nil {
		return nil, nil, err
	}

	memoFile := path.Join(memoDir, fmt.Sprintf("hd_snap%d.dat", snap))

	if _, err := os.Stat(memoFile); err != nil {
		// File not written yet.
		hds, files, err := readHeadersFromSheet(snap)
		if err != nil { return nil, nil, err }

        f, err := os.Create(memoFile)
        if err != nil { return nil, nil, err }
        defer f.Close()
        binary.Write(f, binary.LittleEndian, hds)

		return hds, files, nil
	} else {
		// File exists: read from it instead.

		f, err := os.Open(memoFile)
        if err != nil { return nil, nil, err }
        defer f.Close()

		n, err := sheetNum(snap)
		if err != nil { return nil, nil, err }
		hds := make([]io.SheetHeader, n)
        binary.Read(f, binary.LittleEndian, hds)

		gtetFmt := os.Getenv("GTET_FMT")
		dir := fmt.Sprintf(gtetFmt, snap)
		files, err := dirContents(dir)
		if err != nil { return nil, nil, err }

		return hds, files, nil
	}

}

func sheetNum(snap int) (int, error) {
	gtetFmt := os.Getenv("GTET_FMT")
	dir := fmt.Sprintf(gtetFmt, snap)
	files, err := dirContents(dir)
	if err != nil { return 0, err }
	return len(files), nil
}

func readHeadersFromSheet(snap int) ([]io.SheetHeader, []string, error) {
	gtetFmt := os.Getenv("GTET_FMT")
	dir := fmt.Sprintf(gtetFmt, snap)
	files, err := dirContents(dir)
	if err != nil { return nil, nil, err }

	hds := make([]io.SheetHeader, len(files))
	for i := range files {
		err = io.ReadSheetHeaderAt(files[i], &hds[i])
		if err != nil { return nil, nil, err }
	}
	return hds, files, nil
}

func createHalos(
	snap int, hd *io.SheetHeader, ids []int, p *Params,
) ([]los.HaloProfiles, error) {
	// Read coordinates, radii and IDs of halos.
	rockstarDir := os.Getenv("GTET_ROCKSTAR_DIR")
	if rockstarDir == "" {
		return nil, fmt.Errorf("$GTET_ROCKSTAR_DIR not set.")
	}

	files, err := dirContents(rockstarDir)
	if err != nil { return nil, err }
	// There are no halos in the 0th snapshot
	file := files[snap - 1]

	rids, vals, err := halo.ReadRockstarVals(
		file, &hd.Cosmo, halo.X, halo.Y, halo.Z, halo.M200b,
	)
	if err != nil { return nil, err }

	xs, ys, zs, ms := vals[0], vals[1], vals[2], vals[3]
	halo.R200m.Radius(&hd.Cosmo, ms, ms)
	rs := ms

	// Initialize halos.
	halos := make([]los.HaloProfiles, len(ids))
	seenIDs := make(map[int]bool)
	for i, id := range ids {
		idx := -1
		// TODO: Sort first if the ID count is large enough.
		for j, rid := range rids {
			if rid == id {
				idx = j
				break
			}
		}

		if idx == -1 {
			return nil, fmt.Errorf("Halo ID %d not in halo catalog.", id)
		}

		origin := &geom.Vec{
			float32(xs[idx]), float32(ys[idx]), float32(zs[idx]),
		}

		// If we've already seen a halo once, randomize its orientation.
		if seenIDs[id] {
			halos[i].Init(
				id, p.Rings, origin, rs[idx] * p.MinMult, rs[idx] * p.MaxMult,
				p.RBins, p.Spokes, hd.TotalWidth, los.Log(true),
				los.Rotate(float32(2 * math.Pi * rand.Float64()),
                    float32(2 * math.Pi * rand.Float64()),
                    float32(2 * math.Pi * rand.Float64())),
			)
		} else {
			seenIDs[id] = true
			halos[i].Init(
				id, p.Rings, origin, rs[idx] * p.MinMult, rs[idx] * p.MaxMult,
				p.RBins, p.Spokes, hd.TotalWidth, los.Log(true),
			)
		}
	}

	return halos, nil
}

func binIntersections(
	hds []io.SheetHeader, halos []los.HaloProfiles,
) [][]*los.HaloProfiles {

	bins := make([][]*los.HaloProfiles, len(hds))
	for i := range hds {
		for hi := range halos {
			if (&halos[hi]).SheetIntersect(&hds[i]) {
				bins[i] = append(bins[i], &halos[hi])
			}
		}
	}
	return bins
}

func calcCoeffs(
	halo *los.HaloProfiles, buf []analyze.RingBuffer, p *Params,
) []float64 {
	for i := range buf {
		buf[i].Clear()
		buf[i].Splashback(halo, i, p.Window, p.Cutoff)
	}
	pxs, pys := analyze.FilterPoints(buf, p.Levels)
	cs, _ := analyze.PennaVolumeFit(pxs, pys, halo, p.Order, p.Order)
	return cs
}

func printCoeffs(ids, snaps []int, coeffs [][]float64) {
	idWidth, snapWidth := 0, 0
	coeffWidths := make([]int, len(coeffs[0]))
	for i := range ids {
		iWidth := len(fmt.Sprintf("%d", ids[i]))
		sWidth := len(fmt.Sprintf("%d", snaps[i]))
		if iWidth > idWidth { idWidth = iWidth }
		if sWidth > snapWidth { snapWidth = sWidth }
		for j := range coeffs[i] {
			width := len(fmt.Sprintf("%.5g", coeffs[i][j]))
			if width > coeffWidths[j] { coeffWidths[j] = width }
		}
	}

	rowFmt := fmt.Sprintf("%%%dd %%%dd", idWidth, snapWidth)
	coefFmts := make([]string, len(coeffs[0]))
	for i := range coefFmts {
		rowFmt += fmt.Sprintf(" %%%d.5g", coeffWidths[i])
	}
	rowFmt += "\n"
	for i := range ids {
		args := append([]interface{}{ids[i], snaps[i]}, intr(coeffs[i])...)
		fmt.Printf(rowFmt, args...)
	}
}

func intr(xs []float64) []interface{} {
	is := make([]interface{}, len(xs))
	for i, x := range xs { is[i] = x }
	return is
}

func binBySnap(snaps, ids []int) (snapBins, idxBins map[int][]int) {
	snapBins = make(map[int][]int)
	idxBins = make(map[int][]int)
	for i, snap := range snaps {
		id := ids[i]
		snapBins[snap] = append(snapBins[snap], id)
		idxBins[snap] = append(idxBins[snap], i)
	}
	return snapBins, idxBins
}

func dirContents(dir string) ([]string, error) {
	infos, err := ioutil.ReadDir(dir)
	if err != nil { return nil, err }

	files := make([]string, len(infos))
	for i := range infos {
		files[i] = path.Join(dir, infos[i].Name())
	}

	return files, nil
}
<commit_msg>Removed logging print statement.<commit_after>package main

import (
	"encoding/binary"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"math/rand"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/phil-mansfield/gotetra/los"
	"github.com/phil-mansfield/gotetra/los/geom"
	"github.com/phil-mansfield/gotetra/los/analyze"
	"github.com/phil-mansfield/gotetra/render/io"
	"github.com/phil-mansfield/gotetra/render/halo"
)

type Params struct {
	// HaloProfiles params
	RBins, Spokes, Rings int
	MaxMult, MinMult float64

	// Splashback params
	Order, Window, Levels int
	Cutoff float64
}

func main() {
	// Parse.
	p := parseCmd()
	ids, snaps, err := parseStdin()
	if err != nil { log.Fatal(err.Error()) }

	if len(ids) == 0 { return }

	// Compute coefficients.
	coeffs := make([][]float64, len(ids))
	snapBins, idxBins := binBySnap(snaps, ids)
	buf := make([]analyze.RingBuffer, p.Rings)
	for i := range buf { buf[i].Init(p.Spokes, p.RBins) }

	var losBuf *los.Buffers
	for snap, snapIDs := range snapBins {
		idxs := idxBins[snap]

		if snap == -1 { continue }

		// Bin halos
		hds, files, err := readHeaders(snap)
		if err != nil { err.Error() }
		if losBuf == nil { losBuf = los.NewBuffers(files[0], &hds[0]) }
		halos, err := createHalos(snap, &hds[0], snapIDs, p)
		if err != nil { log.Fatal(err.Error()) }
		intrBins := binIntersections(hds, halos)

		// Add densities. Done header by header to limit I/O time.
		for i := range hds {
			if len(intrBins[i]) == 0 { continue }
			hdContainer := []io.SheetHeader{hds[i]}
			fileContainer := []string{files[i]}
			los.LoadPtrDensities(
				intrBins[i], hdContainer, fileContainer, losBuf,
			)
		}

		// Calculate Penna coefficients.
		for i := range halos {
			coeffs[idxs[i]] = calcCoeffs(&halos[i], buf, p)
		}

	}
	printCoeffs(ids, snaps, coeffs)
}

func parseCmd() *Params {
	// Parse command line.
	p := &Params{}
	flag.IntVar(&p.RBins, "RBins", 256,
		"Number of radial bins used per LoS.")
	flag.IntVar(&p.Spokes, "Spokes", 1024,
		"Number of LoS's used per ring.")
	flag.IntVar(&p.Rings, "Rings", 10,
		"Number of rings used per halo. 3, 4, 6, and 10 rings are\n" +
			"guaranteed to be uniformly spaced.")
	flag.Float64Var(&p.MaxMult, "MaxMult", 3,
		"Ending radius of LoSs as a multiple of R_200m.")
	flag.Float64Var(&p.MinMult, "MinMult", 0.5,
		"Starting radius of LoSs as a multiple of R_200m.")
	flag.IntVar(&p.Order, "Order", 5,
		"Order of the shell fitting function.")
	flag.IntVar(&p.Window, "Window", 121,
		"Number of bins within smoothign window. Must be odd.")
	flag.IntVar(&p.Levels, "Levels", 4,
		"The number of recurve max-finding levels used by the 2D edge finder.")
	flag.Float64Var(&p.Cutoff, "Cutoff", 0.0,
		"The shallowest slope that can be considered a splashback point.")
	flag.Parse()
	return p
}

func parseStdin() (ids, snaps []int, err error) {
	ids, snaps = []int{}, []int{}
	lines, err := stdinLines()
	if err != nil { return nil, nil, err }
	for i, line := range lines {
		rawTokens := strings.Split(line, " ")
		tokens := make([]string, 0, len(rawTokens))
		for _, tok := range rawTokens {
			if len(tok) != 0 { tokens = append(tokens, tok) }
		}

		var (
			id, snap int
			err error
		)
		switch len(tokens) {
		case 0:
			continue
		case 2:
			id, err = strconv.Atoi(tokens[0])
			if err != nil {
				return nil, nil, fmt.Errorf(
					"One line %d of stdin, %s does not parse as an int.",
					i + 1, tokens[0],
				)
			}
			snap, err = strconv.Atoi(tokens[1])
			if err != nil {
				return nil, nil, fmt.Errorf(
					"One line %d of stdin, %s does not parse as an int.",
					i + 1, tokens[1],
				)
			}
		case 1:
			if tokens[0] == "" { continue }
			fallthrough
		default:
			return nil, nil, fmt.Errorf(
				"Line %d of stdin has %d tokens, but 2 are required.",
				i + 1, len(tokens),
			)
		}

		ids = append(ids, id)
		snaps = append(snaps, snap)
	}

	return ids, snaps, nil
}

func stdinLines() ([]string, error) {
		bs, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		return nil, fmt.Errorf(
			"Error reading stdin: %s.", err.Error(),
		)
	}

	text := string(bs)
	return strings.Split(text, "\n"), nil
}

func readHeaders(snap int) ([]io.SheetHeader, []string, error) {
	memoDir := os.Getenv("GTET_MEMO_DIR")
	if memoDir == "" {
		// You don't want to memoize? Fine. Deal with the consequences.
		return readHeadersFromSheet(snap)
	}
	if _, err := os.Stat(memoDir); err != nil {
		return nil, nil, err
	}

	memoFile := path.Join(memoDir, fmt.Sprintf("hd_snap%d.dat", snap))

	if _, err := os.Stat(memoFile); err != nil {
		// File not written yet.
		hds, files, err := readHeadersFromSheet(snap)
		if err != nil { return nil, nil, err }

        f, err := os.Create(memoFile)
        if err != nil { return nil, nil, err }
        defer f.Close()
        binary.Write(f, binary.LittleEndian, hds)

		return hds, files, nil
	} else {
		// File exists: read from it instead.

		f, err := os.Open(memoFile)
        if err != nil { return nil, nil, err }
        defer f.Close()

		n, err := sheetNum(snap)
		if err != nil { return nil, nil, err }
		hds := make([]io.SheetHeader, n)
        binary.Read(f, binary.LittleEndian, hds)

		gtetFmt := os.Getenv("GTET_FMT")
		dir := fmt.Sprintf(gtetFmt, snap)
		files, err := dirContents(dir)
		if err != nil { return nil, nil, err }

		return hds, files, nil
	}

}

func sheetNum(snap int) (int, error) {
	gtetFmt := os.Getenv("GTET_FMT")
	dir := fmt.Sprintf(gtetFmt, snap)
	files, err := dirContents(dir)
	if err != nil { return 0, err }
	return len(files), nil
}

func readHeadersFromSheet(snap int) ([]io.SheetHeader, []string, error) {
	gtetFmt := os.Getenv("GTET_FMT")
	dir := fmt.Sprintf(gtetFmt, snap)
	files, err := dirContents(dir)
	if err != nil { return nil, nil, err }

	hds := make([]io.SheetHeader, len(files))
	for i := range files {
		err = io.ReadSheetHeaderAt(files[i], &hds[i])
		if err != nil { return nil, nil, err }
	}
	return hds, files, nil
}

func createHalos(
	snap int, hd *io.SheetHeader, ids []int, p *Params,
) ([]los.HaloProfiles, error) {
	// Read coordinates, radii and IDs of halos.
	rockstarDir := os.Getenv("GTET_ROCKSTAR_DIR")
	if rockstarDir == "" {
		return nil, fmt.Errorf("$GTET_ROCKSTAR_DIR not set.")
	}

	files, err := dirContents(rockstarDir)
	if err != nil { return nil, err }
	// There are no halos in the 0th snapshot
	file := files[snap - 1]

	rids, vals, err := halo.ReadRockstarVals(
		file, &hd.Cosmo, halo.X, halo.Y, halo.Z, halo.M200b,
	)
	if err != nil { return nil, err }

	xs, ys, zs, ms := vals[0], vals[1], vals[2], vals[3]
	halo.R200m.Radius(&hd.Cosmo, ms, ms)
	rs := ms

	// Initialize halos.
	halos := make([]los.HaloProfiles, len(ids))
	seenIDs := make(map[int]bool)
	for i, id := range ids {
		idx := -1
		// TODO: Sort first if the ID count is large enough.
		for j, rid := range rids {
			if rid == id {
				idx = j
				break
			}
		}

		if idx == -1 {
			return nil, fmt.Errorf("Halo ID %d not in halo catalog.", id)
		}

		origin := &geom.Vec{
			float32(xs[idx]), float32(ys[idx]), float32(zs[idx]),
		}

		// If we've already seen a halo once, randomize its orientation.
		if seenIDs[id] {
			halos[i].Init(
				id, p.Rings, origin, rs[idx] * p.MinMult, rs[idx] * p.MaxMult,
				p.RBins, p.Spokes, hd.TotalWidth, los.Log(true),
				los.Rotate(float32(2 * math.Pi * rand.Float64()),
                    float32(2 * math.Pi * rand.Float64()),
                    float32(2 * math.Pi * rand.Float64())),
			)
		} else {
			seenIDs[id] = true
			halos[i].Init(
				id, p.Rings, origin, rs[idx] * p.MinMult, rs[idx] * p.MaxMult,
				p.RBins, p.Spokes, hd.TotalWidth, los.Log(true),
			)
		}
	}

	return halos, nil
}

func binIntersections(
	hds []io.SheetHeader, halos []los.HaloProfiles,
) [][]*los.HaloProfiles {

	bins := make([][]*los.HaloProfiles, len(hds))
	for i := range hds {
		for hi := range halos {
			if (&halos[hi]).SheetIntersect(&hds[i]) {
				bins[i] = append(bins[i], &halos[hi])
			}
		}
	}
	return bins
}

func calcCoeffs(
	halo *los.HaloProfiles, buf []analyze.RingBuffer, p *Params,
) []float64 {
	for i := range buf {
		buf[i].Clear()
		buf[i].Splashback(halo, i, p.Window, p.Cutoff)
	}
	pxs, pys := analyze.FilterPoints(buf, p.Levels)
	cs, _ := analyze.PennaVolumeFit(pxs, pys, halo, p.Order, p.Order)
	return cs
}

func printCoeffs(ids, snaps []int, coeffs [][]float64) {
	idWidth, snapWidth := 0, 0
	coeffWidths := make([]int, len(coeffs[0]))
	for i := range ids {
		iWidth := len(fmt.Sprintf("%d", ids[i]))
		sWidth := len(fmt.Sprintf("%d", snaps[i]))
		if iWidth > idWidth { idWidth = iWidth }
		if sWidth > snapWidth { snapWidth = sWidth }
		for j := range coeffs[i] {
			width := len(fmt.Sprintf("%.5g", coeffs[i][j]))
			if width > coeffWidths[j] { coeffWidths[j] = width }
		}
	}

	rowFmt := fmt.Sprintf("%%%dd %%%dd", idWidth, snapWidth)
	coefFmts := make([]string, len(coeffs[0]))
	for i := range coefFmts {
		rowFmt += fmt.Sprintf(" %%%d.5g", coeffWidths[i])
	}
	rowFmt += "\n"
	for i := range ids {
		args := append([]interface{}{ids[i], snaps[i]}, intr(coeffs[i])...)
		fmt.Printf(rowFmt, args...)
	}
}

func intr(xs []float64) []interface{} {
	is := make([]interface{}, len(xs))
	for i, x := range xs { is[i] = x }
	return is
}

func binBySnap(snaps, ids []int) (snapBins, idxBins map[int][]int) {
	snapBins = make(map[int][]int)
	idxBins = make(map[int][]int)
	for i, snap := range snaps {
		id := ids[i]
		snapBins[snap] = append(snapBins[snap], id)
		idxBins[snap] = append(idxBins[snap], i)
	}
	return snapBins, idxBins
}

func dirContents(dir string) ([]string, error) {
	infos, err := ioutil.ReadDir(dir)
	if err != nil { return nil, err }

	files := make([]string, len(infos))
	for i := range infos {
		files[i] = path.Join(dir, infos[i].Name())
	}

	return files, nil
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc. All Rights Reserved.
// Copyright 2016 Palm Stone Games, 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 (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"log"
	"sync"
	"time"

	"github.com/boltdb/bolt"
	"github.com/pkg/errors"
	"github.com/xenolf/lego/acme"
	"github.com/xenolf/lego/providers/dns/cloudflare"
	"github.com/xenolf/lego/providers/dns/digitalocean"
	"github.com/xenolf/lego/providers/dns/dnsimple"
	"github.com/xenolf/lego/providers/dns/dnsmadeeasy"
	"github.com/xenolf/lego/providers/dns/dyn"
	"github.com/xenolf/lego/providers/dns/gandi"
	"github.com/xenolf/lego/providers/dns/googlecloud"
	"github.com/xenolf/lego/providers/dns/namecheap"
	"github.com/xenolf/lego/providers/dns/ovh"
	"github.com/xenolf/lego/providers/dns/pdns"
	"github.com/xenolf/lego/providers/dns/rfc2136"
	"github.com/xenolf/lego/providers/dns/route53"
	"github.com/xenolf/lego/providers/dns/vultr"
	"encoding/json"
)

type CertProcessor struct {
	certSecretPrefix string
	acmeURL          string
	db               *bolt.DB
	Lock             sync.Mutex
}

func NewCertProcessor(acmeURL string, certSecretPrefix string, db *bolt.DB) *CertProcessor {
	return &CertProcessor{acmeURL: acmeURL,
		certSecretPrefix: certSecretPrefix,
		db:               db,
	}
}

func (p *CertProcessor) newACMEClient(acmeUser acme.User, provider string) (*acme.Client, error) {
	acmeClient, err := acme.NewClient(p.acmeURL, acmeUser, acme.RSA2048)
	if err != nil {
		return nil, errors.Wrap(err, "Error while generating acme client")
	}

	initDNSProvider := func(p acme.ChallengeProvider, err error) (*acme.Client, error) {
		if err != nil {
			return nil, errors.Wrapf(err, "Error while initializing provider %v", provider)
		}

		if err := acmeClient.SetChallengeProvider(acme.DNS01, p); err != nil {
			return nil, errors.Wrap(err, "Error while setting cloudflare challenge provider")
		}

		acmeClient.ExcludeChallenges([]acme.Challenge{acme.HTTP01, acme.TLSSNI01})
		return acmeClient, nil
	}

	switch provider {
	case "http":
		acmeClient.SetHTTPAddress("8080")
		acmeClient.ExcludeChallenges([]acme.Challenge{acme.DNS01, acme.TLSSNI01})
		return acmeClient, nil
	case "tls":
		acmeClient.SetTLSAddress("8081")
		acmeClient.ExcludeChallenges([]acme.Challenge{acme.HTTP01, acme.DNS01})
		return acmeClient, nil
	case "cloudflare":
		return initDNSProvider(cloudflare.NewDNSProvider())
	case "digitalocean":
		return initDNSProvider(digitalocean.NewDNSProvider())
	case "dnsimple":
		return initDNSProvider(dnsimple.NewDNSProvider())
	case "dnsmadeeasy":
		return initDNSProvider(dnsmadeeasy.NewDNSProvider())
	case "dyn":
		return initDNSProvider(dyn.NewDNSProvider())
	case "gandi":
		return initDNSProvider(gandi.NewDNSProvider())
	case "googlecloud":
		return initDNSProvider(googlecloud.NewDNSProvider())
	case "namecheap":
		return initDNSProvider(namecheap.NewDNSProvider())
	case "ovh":
		return initDNSProvider(ovh.NewDNSProvider())
	case "pdns":
		return initDNSProvider(pdns.NewDNSProvider())
	case "rfc2136":
		return initDNSProvider(rfc2136.NewDNSProvider())
	case "route53":
		return initDNSProvider(route53.NewDNSProvider())
	case "vultr":
		return initDNSProvider(vultr.NewDNSProvider())
	default:
		return nil, errors.Errorf("Unknown provider %v", provider)
	}
}

func (p *CertProcessor) syncCertificates(verbose bool) error {
	p.Lock.Lock()
	defer p.Lock.Unlock()

	certificates, err := getCertificates()
	if err != nil {
		return errors.Wrap(err, "Error while fetching certificate list")
	}

	var wg sync.WaitGroup
	for _, cert := range certificates {
		wg.Add(1)
		go func(cert Certificate) {
			defer wg.Done()
			err := p.processCertificate(&cert.Spec)
			if err != nil {
				log.Printf("Error while processing certificate during sync: %v", err)
			}
		}(cert)
	}
	wg.Wait()
	return nil
}

func (p *CertProcessor) watchKubernetesEvents(wg *sync.WaitGroup, doneChan <-chan struct{}) {
	events, watchErrs := monitorCertificateEvents()
	go func() {
		for {
			select {
			case event := <-events:
				err := p.processCertificateEvent(event)
				if err != nil {
					log.Printf("Error while processing certificate event: %v", err)
				}
			case err := <-watchErrs:
				log.Printf("Error while watching kubernetes events: %v", err)
			case <-doneChan:
				wg.Done()
				log.Println("Stopped certificate event watcher.")
				return
			}
		}
	}()
}

func (p *CertProcessor) refreshCertificates(syncInterval time.Duration, wg *sync.WaitGroup, doneChan <-chan struct{}) {
	go func() {
		for {
			select {
			case <-time.After(syncInterval):
				err := p.syncCertificates(false)
				if err != nil {
					log.Printf("Error while synchronizing certificates during refresh: %v", err)
				}
			case <-doneChan:
				wg.Done()
				log.Println("Stopped refresh loop.")
				return
			}
		}
	}()
}

func (p *CertProcessor) processCertificateEvent(c CertificateEvent) error {
	p.Lock.Lock()
	defer p.Lock.Unlock()
	switch {
	case c.Type == "ADDED":
		return p.processCertificate(&c.Object.Spec)
	case c.Type == "DELETED":
		return p.deleteCertificate(&c.Object.Spec)
	}
	return nil
}

func (p *CertProcessor) processCertificate(certSpec *CertificateSpec) error {
	var (
		acmeUserInfo    ACMEUserData
		acmeCertDetails ACMECertDetails
		acmeCert        ACMECertData
		acmeClient 	*acme.Client
	)

	// Fetch current certificate data from k8s
	s, err := getSecret(p.certSecretPrefix + certSpec.Domain)
	if err != nil {
		return errors.Wrapf(err, "Error while fetching certificate acme data for domain %v", certSpec.Domain)
	}

	// If a cert exists, check its expiry
	if s != nil {
		acmeCert, err = NewACMECertDataFromSecret(s)
		if err != nil {
			return errors.Wrapf(err, "Error while decoding acme certificate from secret for existing domain %v", certSpec.Domain)
		}

		// Decode cert
		pemBlock, _ := pem.Decode(acmeCert.Cert)
		parsedCert, err := x509.ParseCertificate(pemBlock.Bytes)
		if err != nil {
			return errors.Wrapf(err, "Error while decoding x509 encoded certificate for existing domain %v", certSpec.Domain)
		}

		// If certificate expires in more than a week, don't renew
		if parsedCert.NotAfter.After(time.Now().Add(time.Hour * 24 * 7)) {
			return nil
		}

		log.Printf("[%v] Expiry for cert is in less than a week (%v), attempting renewal", certSpec.Domain, parsedCert.NotAfter.String())
	}

	// Fetch acme user data and cert details from bolt
	var userInfoRaw, certDetailsRaw []byte
	err = p.db.View(func(tx *bolt.Tx) error {
		userInfoRaw = tx.Bucket([]byte("user-info")).Get([]byte(certSpec.Domain))
		certDetailsRaw = tx.Bucket([]byte("cert-details")).Get([]byte(certSpec.Domain))
		return nil
	})

	if err != nil {
		return errors.Wrapf(err, "Error while running bolt view transaction for domain %v", certSpec.Domain)
	}

	// Handle user information
	if userInfoRaw != nil { // Use existing user
		if err := json.Unmarshal(userInfoRaw, &acmeUserInfo); err != nil {
			return errors.Wrapf(err, "Error while unmarshalling user info for %v", certSpec.Domain)
		}

		acmeClient, err = p.newACMEClient(&acmeUserInfo, certSpec.Provider)
		if err != nil {
			return errors.Wrapf(err, "Error while creating ACME client for %v", certSpec.Domain)
		}
	} else { // Generate a new ACME user
		userKey, err := rsa.GenerateKey(rand.Reader, 2048)
		if err != nil {
			return errors.Wrapf(err, "Error while generating rsa key for new user for domain %v", certSpec.Domain)
		}

		acmeUserInfo.Email = certSpec.Email
		acmeUserInfo.Key = pem.EncodeToMemory(&pem.Block{
			Type:  "RSA PRIVATE KEY",
			Bytes: x509.MarshalPKCS1PrivateKey(userKey),
		})

		acmeClient, err = p.newACMEClient(&acmeUserInfo, certSpec.Provider)
		if err != nil {
			return errors.Wrapf(err, "Error while creating ACME client for %v", certSpec.Domain)
		}

		// Register
		acmeUserInfo.Registration, err = acmeClient.Register()
		if err != nil {
			return errors.Wrapf(err, "Error while registering user for new domain %v", certSpec.Domain)
		}

		// Agree to TOS
		if err := acmeClient.AgreeToTOS(); err != nil {
			return errors.Wrapf(err, "Error while agreeing to acme TOS for new domain %v", certSpec.Domain)
		}
	}

	// If we have cert details stored, do a renewal, otherwise, obtain from scratch
	if certDetailsRaw == nil || acmeCert.DomainName == "" {
		acmeCert.DomainName = certSpec.Domain

		// Obtain a cert
		certRes, errs := acmeClient.ObtainCertificate([]string{certSpec.Domain}, true, nil)
		if errs[certSpec.Domain] != nil {
			return errors.Wrapf(errs[certSpec.Domain], "Error while obtaining certificate for new domain %v", certSpec.Domain)
		}

		// fill in data
		acmeCert.Cert = certRes.Certificate
		acmeCert.PrivateKey = certRes.PrivateKey
		acmeCertDetails = NewACMECertDetailsFromResource(certRes)
	} else {
		if err := json.Unmarshal(certDetailsRaw, &acmeCertDetails); err != nil {
			return errors.Wrapf(err, "Error while unmarshalling cert details for existing domain %v", certSpec.Domain)
		}

		// Initialize ACME client
		acmeClient, err := p.newACMEClient(&acmeUserInfo, certSpec.Provider)
		if err != nil {
			return errors.Wrapf(err, "Error while creating ACME client for existing domain %v", certSpec.Domain)
		}

		// Fill in cert resource
		certRes := acmeCertDetails.ToCertResource()
		certRes.Certificate = acmeCert.Cert
		certRes.PrivateKey = acmeCert.PrivateKey

		certRes, err = acmeClient.RenewCertificate(certRes, true)
		if err != nil {
			return errors.Wrapf(err, "Error while renewing certificate for existing domain %v", certSpec.Domain)
		}

		// Fill in details
		acmeCert.Cert = certRes.Certificate
		acmeCert.PrivateKey = certRes.PrivateKey
		acmeCertDetails = NewACMECertDetailsFromResource(certRes)
	}

	// Serialize acmeCertDetails and acmeUserInfo
	certDetailsRaw, err = json.Marshal(&acmeCertDetails)
	if err != nil {
		return errors.Wrapf(err, "Error while marshalling cert details for domain %v", certSpec.Domain)
	}

	userInfoRaw, err = json.Marshal(&acmeUserInfo)
	if err != nil {
		return errors.Wrapf(err, "Error while marshalling user info for domain %v", certSpec.Domain)
	}

	// Save cert details and user info to bolt
	err = p.db.Update(func(tx *bolt.Tx) error {
		key := []byte(certSpec.Domain)
		tx.Bucket([]byte("user-info")).Put(key, userInfoRaw)
		tx.Bucket([]byte("cert-details")).Put(key, certDetailsRaw)
		return nil
	})
	if err != nil {
		return errors.Wrapf(err, "Error while saving data to bolt for domain %v", certSpec.Domain)
	}

	// Convert cert data to k8s secret
	isUpdate := s != nil
	s, err = acmeCert.ToSecret(p.certSecretPrefix)
	if err != nil {
		return errors.Wrapf(err, "Error while creating secret for ACME certificate for domain %v", certSpec.Domain)
	}

	// Save the k8s secret
	if err := saveSecret(s, isUpdate); err != nil {
		return errors.Wrapf(err, "Error while saving secret for domain %v", certSpec.Domain)
	}

	return nil
}

func (p *CertProcessor) deleteCertificate(certSpec *CertificateSpec) error {
	return deleteSecret(p.certSecretPrefix + certSpec.Domain)
}
<commit_msg>Only allow a single http/tls validation at once<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
// Copyright 2016 Palm Stone Games, 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 (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"log"
	"sync"
	"time"
	"encoding/json"

	"github.com/boltdb/bolt"
	"github.com/pkg/errors"
	"github.com/xenolf/lego/acme"
	"github.com/xenolf/lego/providers/dns/cloudflare"
	"github.com/xenolf/lego/providers/dns/digitalocean"
	"github.com/xenolf/lego/providers/dns/dnsimple"
	"github.com/xenolf/lego/providers/dns/dnsmadeeasy"
	"github.com/xenolf/lego/providers/dns/dyn"
	"github.com/xenolf/lego/providers/dns/gandi"
	"github.com/xenolf/lego/providers/dns/googlecloud"
	"github.com/xenolf/lego/providers/dns/namecheap"
	"github.com/xenolf/lego/providers/dns/ovh"
	"github.com/xenolf/lego/providers/dns/pdns"
	"github.com/xenolf/lego/providers/dns/rfc2136"
	"github.com/xenolf/lego/providers/dns/route53"
	"github.com/xenolf/lego/providers/dns/vultr"
)

type CertProcessor struct {
	certSecretPrefix string
	acmeURL          string
	db               *bolt.DB
	Lock             sync.Mutex
	HTTPLock         sync.Mutex
	TLSLock          sync.Mutex
}

func NewCertProcessor(acmeURL string, certSecretPrefix string, db *bolt.DB) *CertProcessor {
	return &CertProcessor{acmeURL: acmeURL,
		certSecretPrefix: certSecretPrefix,
		db:               db,
	}
}

func (p *CertProcessor) newACMEClient(acmeUser acme.User, provider string) (*acme.Client, *sync.Mutex, error) {
	acmeClient, err := acme.NewClient(p.acmeURL, acmeUser, acme.RSA2048)
	if err != nil {
		return nil, nil, errors.Wrap(err, "Error while generating acme client")
	}

	initDNSProvider := func(p acme.ChallengeProvider, err error) (*acme.Client, *sync.Mutex, error) {
		if err != nil {
			return nil, nil, errors.Wrapf(err, "Error while initializing provider %v", provider)
		}

		if err := acmeClient.SetChallengeProvider(acme.DNS01, p); err != nil {
			return nil, nil, errors.Wrap(err, "Error while setting cloudflare challenge provider")
		}

		acmeClient.ExcludeChallenges([]acme.Challenge{acme.HTTP01, acme.TLSSNI01})
		return acmeClient, nil, nil
	}

	switch provider {
	case "http":
		acmeClient.SetHTTPAddress("8080")
		acmeClient.ExcludeChallenges([]acme.Challenge{acme.DNS01, acme.TLSSNI01})
		return acmeClient, &p.HTTPLock, nil
	case "tls":
		acmeClient.SetTLSAddress("8081")
		acmeClient.ExcludeChallenges([]acme.Challenge{acme.HTTP01, acme.DNS01})
		return acmeClient, &p.TLSLock, nil
	case "cloudflare":
		return initDNSProvider(cloudflare.NewDNSProvider())
	case "digitalocean":
		return initDNSProvider(digitalocean.NewDNSProvider())
	case "dnsimple":
		return initDNSProvider(dnsimple.NewDNSProvider())
	case "dnsmadeeasy":
		return initDNSProvider(dnsmadeeasy.NewDNSProvider())
	case "dyn":
		return initDNSProvider(dyn.NewDNSProvider())
	case "gandi":
		return initDNSProvider(gandi.NewDNSProvider())
	case "googlecloud":
		return initDNSProvider(googlecloud.NewDNSProvider())
	case "namecheap":
		return initDNSProvider(namecheap.NewDNSProvider())
	case "ovh":
		return initDNSProvider(ovh.NewDNSProvider())
	case "pdns":
		return initDNSProvider(pdns.NewDNSProvider())
	case "rfc2136":
		return initDNSProvider(rfc2136.NewDNSProvider())
	case "route53":
		return initDNSProvider(route53.NewDNSProvider())
	case "vultr":
		return initDNSProvider(vultr.NewDNSProvider())
	default:
		return nil, nil, errors.Errorf("Unknown provider %v", provider)
	}
}

func (p *CertProcessor) syncCertificates(verbose bool) error {
	p.Lock.Lock()
	defer p.Lock.Unlock()

	certificates, err := getCertificates()
	if err != nil {
		return errors.Wrap(err, "Error while fetching certificate list")
	}

	var wg sync.WaitGroup
	for _, cert := range certificates {
		wg.Add(1)
		go func(cert Certificate) {
			defer wg.Done()
			err := p.processCertificate(&cert.Spec)
			if err != nil {
				log.Printf("Error while processing certificate during sync: %v", err)
			}
		}(cert)
	}
	wg.Wait()
	return nil
}

func (p *CertProcessor) watchKubernetesEvents(wg *sync.WaitGroup, doneChan <-chan struct{}) {
	events, watchErrs := monitorCertificateEvents()
	go func() {
		for {
			select {
			case event := <-events:
				err := p.processCertificateEvent(event)
				if err != nil {
					log.Printf("Error while processing certificate event: %v", err)
				}
			case err := <-watchErrs:
				log.Printf("Error while watching kubernetes events: %v", err)
			case <-doneChan:
				wg.Done()
				log.Println("Stopped certificate event watcher.")
				return
			}
		}
	}()
}

func (p *CertProcessor) refreshCertificates(syncInterval time.Duration, wg *sync.WaitGroup, doneChan <-chan struct{}) {
	go func() {
		for {
			select {
			case <-time.After(syncInterval):
				err := p.syncCertificates(false)
				if err != nil {
					log.Printf("Error while synchronizing certificates during refresh: %v", err)
				}
			case <-doneChan:
				wg.Done()
				log.Println("Stopped refresh loop.")
				return
			}
		}
	}()
}

func (p *CertProcessor) processCertificateEvent(c CertificateEvent) error {
	p.Lock.Lock()
	defer p.Lock.Unlock()
	switch {
	case c.Type == "ADDED":
		return p.processCertificate(&c.Object.Spec)
	case c.Type == "DELETED":
		return p.deleteCertificate(&c.Object.Spec)
	}
	return nil
}

func (p *CertProcessor) processCertificate(certSpec *CertificateSpec) error {
	var (
		acmeUserInfo    ACMEUserData
		acmeCertDetails ACMECertDetails
		acmeCert        ACMECertData
		acmeClient      *acme.Client
		acmeClientMutex *sync.Mutex
	)

	// Fetch current certificate data from k8s
	s, err := getSecret(p.certSecretPrefix + certSpec.Domain)
	if err != nil {
		return errors.Wrapf(err, "Error while fetching certificate acme data for domain %v", certSpec.Domain)
	}

	// If a cert exists, check its expiry
	if s != nil {
		acmeCert, err = NewACMECertDataFromSecret(s)
		if err != nil {
			return errors.Wrapf(err, "Error while decoding acme certificate from secret for existing domain %v", certSpec.Domain)
		}

		// Decode cert
		pemBlock, _ := pem.Decode(acmeCert.Cert)
		parsedCert, err := x509.ParseCertificate(pemBlock.Bytes)
		if err != nil {
			return errors.Wrapf(err, "Error while decoding x509 encoded certificate for existing domain %v", certSpec.Domain)
		}

		// If certificate expires in more than a week, don't renew
		if parsedCert.NotAfter.After(time.Now().Add(time.Hour * 24 * 7)) {
			return nil
		}

		log.Printf("[%v] Expiry for cert is in less than a week (%v), attempting renewal", certSpec.Domain, parsedCert.NotAfter.String())
	}

	// Fetch acme user data and cert details from bolt
	var userInfoRaw, certDetailsRaw []byte
	err = p.db.View(func(tx *bolt.Tx) error {
		userInfoRaw = tx.Bucket([]byte("user-info")).Get([]byte(certSpec.Domain))
		certDetailsRaw = tx.Bucket([]byte("cert-details")).Get([]byte(certSpec.Domain))
		return nil
	})

	if err != nil {
		return errors.Wrapf(err, "Error while running bolt view transaction for domain %v", certSpec.Domain)
	}

	// Handle user information
	if userInfoRaw != nil { // Use existing user
		if err := json.Unmarshal(userInfoRaw, &acmeUserInfo); err != nil {
			return errors.Wrapf(err, "Error while unmarshalling user info for %v", certSpec.Domain)
		}

		acmeClient, acmeClientMutex, err = p.newACMEClient(&acmeUserInfo, certSpec.Provider)
		if err != nil {
			return errors.Wrapf(err, "Error while creating ACME client for %v", certSpec.Domain)
		}

		// Some acme providers require locking, if the mutex is specified, lock it
		if acmeClientMutex != nil {
			acmeClientMutex.Lock()
			defer acmeClientMutex.Lock()
		}
	} else { // Generate a new ACME user
		userKey, err := rsa.GenerateKey(rand.Reader, 2048)
		if err != nil {
			return errors.Wrapf(err, "Error while generating rsa key for new user for domain %v", certSpec.Domain)
		}

		acmeUserInfo.Email = certSpec.Email
		acmeUserInfo.Key = pem.EncodeToMemory(&pem.Block{
			Type:  "RSA PRIVATE KEY",
			Bytes: x509.MarshalPKCS1PrivateKey(userKey),
		})

		acmeClient, acmeClientMutex, err = p.newACMEClient(&acmeUserInfo, certSpec.Provider)
		if err != nil {
			return errors.Wrapf(err, "Error while creating ACME client for %v", certSpec.Domain)
		}

		// Some acme providers require locking, if the mutex is specified, lock it
		if acmeClientMutex != nil {
			acmeClientMutex.Lock()
			defer acmeClientMutex.Lock()
		}

		// Register
		acmeUserInfo.Registration, err = acmeClient.Register()
		if err != nil {
			return errors.Wrapf(err, "Error while registering user for new domain %v", certSpec.Domain)
		}

		// Agree to TOS
		if err := acmeClient.AgreeToTOS(); err != nil {
			return errors.Wrapf(err, "Error while agreeing to acme TOS for new domain %v", certSpec.Domain)
		}
	}

	// If we have cert details stored, do a renewal, otherwise, obtain from scratch
	if certDetailsRaw == nil || acmeCert.DomainName == "" {
		acmeCert.DomainName = certSpec.Domain

		// Obtain a cert
		certRes, errs := acmeClient.ObtainCertificate([]string{certSpec.Domain}, true, nil)
		if errs[certSpec.Domain] != nil {
			return errors.Wrapf(errs[certSpec.Domain], "Error while obtaining certificate for new domain %v", certSpec.Domain)
		}

		// fill in data
		acmeCert.Cert = certRes.Certificate
		acmeCert.PrivateKey = certRes.PrivateKey
		acmeCertDetails = NewACMECertDetailsFromResource(certRes)
	} else {
		if err := json.Unmarshal(certDetailsRaw, &acmeCertDetails); err != nil {
			return errors.Wrapf(err, "Error while unmarshalling cert details for existing domain %v", certSpec.Domain)
		}

		// Fill in cert resource
		certRes := acmeCertDetails.ToCertResource()
		certRes.Certificate = acmeCert.Cert
		certRes.PrivateKey = acmeCert.PrivateKey

		certRes, err = acmeClient.RenewCertificate(certRes, true)
		if err != nil {
			return errors.Wrapf(err, "Error while renewing certificate for existing domain %v", certSpec.Domain)
		}

		// Fill in details
		acmeCert.Cert = certRes.Certificate
		acmeCert.PrivateKey = certRes.PrivateKey
		acmeCertDetails = NewACMECertDetailsFromResource(certRes)
	}

	// Serialize acmeCertDetails and acmeUserInfo
	certDetailsRaw, err = json.Marshal(&acmeCertDetails)
	if err != nil {
		return errors.Wrapf(err, "Error while marshalling cert details for domain %v", certSpec.Domain)
	}

	userInfoRaw, err = json.Marshal(&acmeUserInfo)
	if err != nil {
		return errors.Wrapf(err, "Error while marshalling user info for domain %v", certSpec.Domain)
	}

	// Save cert details and user info to bolt
	err = p.db.Update(func(tx *bolt.Tx) error {
		key := []byte(certSpec.Domain)
		tx.Bucket([]byte("user-info")).Put(key, userInfoRaw)
		tx.Bucket([]byte("cert-details")).Put(key, certDetailsRaw)
		return nil
	})
	if err != nil {
		return errors.Wrapf(err, "Error while saving data to bolt for domain %v", certSpec.Domain)
	}

	// Convert cert data to k8s secret
	isUpdate := s != nil
	s, err = acmeCert.ToSecret(p.certSecretPrefix)
	if err != nil {
		return errors.Wrapf(err, "Error while creating secret for ACME certificate for domain %v", certSpec.Domain)
	}

	// Save the k8s secret
	if err := saveSecret(s, isUpdate); err != nil {
		return errors.Wrapf(err, "Error while saving secret for domain %v", certSpec.Domain)
	}

	return nil
}

func (p *CertProcessor) deleteCertificate(certSpec *CertificateSpec) error {
	return deleteSecret(p.certSecretPrefix + certSpec.Domain)
}
<|endoftext|>


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

package protect

import (
	"context"
	"errors"
	"fmt"
	"net"
	"strings"
	"syscall"
)

// Protector provides the ability to bypass a VPN on Android, pre-Lollipop.
type Protector interface {
	// Protect a socket, i.e. exclude it from the VPN.
	// This is needed in order to avoid routing loops for the VPN's own sockets.
	// This is a wrapper for Android's VpnService.protect().
	Protect(socket int32) bool

	// Returns a comma-separated list of the system's configured DNS resolvers,
	// in roughly descending priority order.
	// This is needed because (1) Android Java cannot protect DNS lookups but Go can, and
	// (2) Android Java can determine the list of system DNS resolvers but Go cannot.
	// A comma-separated list is used because Gomobile cannot bind []string.
	GetResolvers() string
}

func makeControl(p Protector) func(string, string, syscall.RawConn) error {
	return func(network, address string, c syscall.RawConn) error {
		return c.Control(func(fd uintptr) {
			if !p.Protect(int32(fd)) {
				panic("Failed to protect socket")
			}
		})
	}
}

// Returns the first IP address that is of the desired family.
func scan(ips []string, wantV4 bool) string {
	for _, ip := range ips {
		parsed := net.ParseIP(ip)
		if parsed == nil {
			// `ip` failed to parse.  Skip it.
			continue
		}
		isV4 := parsed.To4() != nil
		if isV4 == wantV4 {
			return ip
		}
	}
	return ""
}

// Given a slice of IP addresses, and a transport address, return a transport
// address with the IP replaced by the first IP of the same family in `ips`, or
// by the first address of a different family if there are none of the same.
func replaceIP(addr string, ips []string) (string, error) {
	if len(ips) == 0 {
		return "", errors.New("No resolvers available")
	}
	orighost, port, err := net.SplitHostPort(addr)
	if err != nil {
		return "", err
	}
	origip := net.ParseIP(orighost)
	if origip == nil {
		return "", fmt.Errorf("Can't parse resolver IP: %s", orighost)
	}
	isV4 := origip.To4() != nil
	newIP := scan(ips, isV4)
	if newIP == "" {
		// There are no IPs of the desired address family.  Use a different family.
		newIP = ips[0]
	}
	return net.JoinHostPort(newIP, port), nil
}

// MakeDialer creates a new Dialer.  Recipients can safely mutate
// any public field except Control and Resolver, which are both populated.
func MakeDialer(p Protector) *net.Dialer {
	if p == nil {
		return &net.Dialer{}
	}
	d := &net.Dialer{
		Control: makeControl(p),
	}
	resolverDialer := func(ctx context.Context, network, address string) (net.Conn, error) {
		resolvers := strings.Split(p.GetResolvers(), ",")
		newAddress, err := replaceIP(address, resolvers)
		if err != nil {
			return nil, err
		}
		return d.DialContext(ctx, network, newAddress)
	}
	d.Resolver = &net.Resolver{
		PreferGo: true,
		Dial:     resolverDialer,
	}
	return d
}

// MakeListenConfig returns a new ListenConfig that creates protected
// listener sockets.
func MakeListenConfig(p Protector) *net.ListenConfig {
	if p == nil {
		return &net.ListenConfig{}
	}
	return &net.ListenConfig{
		Control: makeControl(p),
	}
}
<commit_msg>Don't panic if protect() returns false<commit_after>// Copyright 2019 The Outline Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package protect

import (
	"context"
	"errors"
	"fmt"
	"net"
	"strings"
	"syscall"

	"github.com/eycorsican/go-tun2socks/common/log"
)

// Protector provides the ability to bypass a VPN on Android, pre-Lollipop.
type Protector interface {
	// Protect a socket, i.e. exclude it from the VPN.
	// This is needed in order to avoid routing loops for the VPN's own sockets.
	// This is a wrapper for Android's VpnService.protect().
	Protect(socket int32) bool

	// Returns a comma-separated list of the system's configured DNS resolvers,
	// in roughly descending priority order.
	// This is needed because (1) Android Java cannot protect DNS lookups but Go can, and
	// (2) Android Java can determine the list of system DNS resolvers but Go cannot.
	// A comma-separated list is used because Gomobile cannot bind []string.
	GetResolvers() string
}

func makeControl(p Protector) func(string, string, syscall.RawConn) error {
	return func(network, address string, c syscall.RawConn) error {
		return c.Control(func(fd uintptr) {
			if !p.Protect(int32(fd)) {
				log.Errorf("Failed to protect a %s socket", network)
			}
		})
	}
}

// Returns the first IP address that is of the desired family.
func scan(ips []string, wantV4 bool) string {
	for _, ip := range ips {
		parsed := net.ParseIP(ip)
		if parsed == nil {
			// `ip` failed to parse.  Skip it.
			continue
		}
		isV4 := parsed.To4() != nil
		if isV4 == wantV4 {
			return ip
		}
	}
	return ""
}

// Given a slice of IP addresses, and a transport address, return a transport
// address with the IP replaced by the first IP of the same family in `ips`, or
// by the first address of a different family if there are none of the same.
func replaceIP(addr string, ips []string) (string, error) {
	if len(ips) == 0 {
		return "", errors.New("No resolvers available")
	}
	orighost, port, err := net.SplitHostPort(addr)
	if err != nil {
		return "", err
	}
	origip := net.ParseIP(orighost)
	if origip == nil {
		return "", fmt.Errorf("Can't parse resolver IP: %s", orighost)
	}
	isV4 := origip.To4() != nil
	newIP := scan(ips, isV4)
	if newIP == "" {
		// There are no IPs of the desired address family.  Use a different family.
		newIP = ips[0]
	}
	return net.JoinHostPort(newIP, port), nil
}

// MakeDialer creates a new Dialer.  Recipients can safely mutate
// any public field except Control and Resolver, which are both populated.
func MakeDialer(p Protector) *net.Dialer {
	if p == nil {
		return &net.Dialer{}
	}
	d := &net.Dialer{
		Control: makeControl(p),
	}
	resolverDialer := func(ctx context.Context, network, address string) (net.Conn, error) {
		resolvers := strings.Split(p.GetResolvers(), ",")
		newAddress, err := replaceIP(address, resolvers)
		if err != nil {
			return nil, err
		}
		return d.DialContext(ctx, network, newAddress)
	}
	d.Resolver = &net.Resolver{
		PreferGo: true,
		Dial:     resolverDialer,
	}
	return d
}

// MakeListenConfig returns a new ListenConfig that creates protected
// listener sockets.
func MakeListenConfig(p Protector) *net.ListenConfig {
	if p == nil {
		return &net.ListenConfig{}
	}
	return &net.ListenConfig{
		Control: makeControl(p),
	}
}
<|endoftext|>


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

package protect

import (
	"context"
	"errors"
	"fmt"
	"net"
	"strings"
	"syscall"

	"github.com/eycorsican/go-tun2socks/common/log"
)

// Protector provides the ability to bypass a VPN on Android, pre-Lollipop.
type Protector interface {
	// Protect a socket, i.e. exclude it from the VPN.
	// This is needed in order to avoid routing loops for the VPN's own sockets.
	// This is a wrapper for Android's VpnService.protect().
	Protect(socket int32) bool

	// Returns a comma-separated list of the system's configured DNS resolvers,
	// in roughly descending priority order.
	// This is needed because (1) Android Java cannot protect DNS lookups but Go can, and
	// (2) Android Java can determine the list of system DNS resolvers but Go cannot.
	// A comma-separated list is used because Gomobile cannot bind []string.
	GetResolvers() string
}

func makeControl(p Protector) func(string, string, syscall.RawConn) error {
	return func(network, address string, c syscall.RawConn) error {
		return c.Control(func(fd uintptr) {
			if !p.Protect(int32(fd)) {
				log.Errorf("Failed to protect a %s socket", network)
			}
		})
	}
}

// Returns the first IP address that is of the desired family.
func scan(ips []string, wantV4 bool) string {
	for _, ip := range ips {
		parsed := net.ParseIP(ip)
		if parsed == nil {
			// `ip` failed to parse.  Skip it.
			continue
		}
		isV4 := parsed.To4() != nil
		if isV4 == wantV4 {
			return ip
		}
	}
	return ""
}

// Given a slice of IP addresses, and a transport address, return a transport
// address with the IP replaced by the first IP of the same family in `ips`, or
// by the first address of a different family if there are none of the same.
func replaceIP(addr string, ips []string) (string, error) {
	if len(ips) == 0 {
		return "", errors.New("No resolvers available")
	}
	orighost, port, err := net.SplitHostPort(addr)
	if err != nil {
		return "", err
	}
	origip := net.ParseIP(orighost)
	if origip == nil {
		return "", fmt.Errorf("Can't parse resolver IP: %s", orighost)
	}
	isV4 := origip.To4() != nil
	newIP := scan(ips, isV4)
	if newIP == "" {
		// There are no IPs of the desired address family.  Use a different family.
		newIP = ips[0]
	}
	return net.JoinHostPort(newIP, port), nil
}

// MakeDialer creates a new Dialer.  Recipients can safely mutate
// any public field except Control and Resolver, which are both populated.
func MakeDialer(p Protector) *net.Dialer {
	if p == nil {
		return &net.Dialer{}
	}
	d := &net.Dialer{
		Control: makeControl(p),
	}
	resolverDialer := func(ctx context.Context, network, address string) (net.Conn, error) {
		resolvers := strings.Split(p.GetResolvers(), ",")
		newAddress, err := replaceIP(address, resolvers)
		if err != nil {
			return nil, err
		}
		return d.DialContext(ctx, network, newAddress)
	}
	d.Resolver = &net.Resolver{
		PreferGo: true,
		Dial:     resolverDialer,
	}
	return d
}

// MakeListenConfig returns a new ListenConfig that creates protected
// listener sockets.
func MakeListenConfig(p Protector) *net.ListenConfig {
	if p == nil {
		return &net.ListenConfig{}
	}
	return &net.ListenConfig{
		Control: makeControl(p),
	}
}
<commit_msg>Add a TODO for future work<commit_after>// Copyright 2019 The Outline Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package protect

import (
	"context"
	"errors"
	"fmt"
	"net"
	"strings"
	"syscall"

	"github.com/eycorsican/go-tun2socks/common/log"
)

// Protector provides the ability to bypass a VPN on Android, pre-Lollipop.
type Protector interface {
	// Protect a socket, i.e. exclude it from the VPN.
	// This is needed in order to avoid routing loops for the VPN's own sockets.
	// This is a wrapper for Android's VpnService.protect().
	Protect(socket int32) bool

	// Returns a comma-separated list of the system's configured DNS resolvers,
	// in roughly descending priority order.
	// This is needed because (1) Android Java cannot protect DNS lookups but Go can, and
	// (2) Android Java can determine the list of system DNS resolvers but Go cannot.
	// A comma-separated list is used because Gomobile cannot bind []string.
	GetResolvers() string
}

func makeControl(p Protector) func(string, string, syscall.RawConn) error {
	return func(network, address string, c syscall.RawConn) error {
		return c.Control(func(fd uintptr) {
			if !p.Protect(int32(fd)) {
				// TODO: Record and report these errors.
				log.Errorf("Failed to protect a %s socket", network)
			}
		})
	}
}

// Returns the first IP address that is of the desired family.
func scan(ips []string, wantV4 bool) string {
	for _, ip := range ips {
		parsed := net.ParseIP(ip)
		if parsed == nil {
			// `ip` failed to parse.  Skip it.
			continue
		}
		isV4 := parsed.To4() != nil
		if isV4 == wantV4 {
			return ip
		}
	}
	return ""
}

// Given a slice of IP addresses, and a transport address, return a transport
// address with the IP replaced by the first IP of the same family in `ips`, or
// by the first address of a different family if there are none of the same.
func replaceIP(addr string, ips []string) (string, error) {
	if len(ips) == 0 {
		return "", errors.New("No resolvers available")
	}
	orighost, port, err := net.SplitHostPort(addr)
	if err != nil {
		return "", err
	}
	origip := net.ParseIP(orighost)
	if origip == nil {
		return "", fmt.Errorf("Can't parse resolver IP: %s", orighost)
	}
	isV4 := origip.To4() != nil
	newIP := scan(ips, isV4)
	if newIP == "" {
		// There are no IPs of the desired address family.  Use a different family.
		newIP = ips[0]
	}
	return net.JoinHostPort(newIP, port), nil
}

// MakeDialer creates a new Dialer.  Recipients can safely mutate
// any public field except Control and Resolver, which are both populated.
func MakeDialer(p Protector) *net.Dialer {
	if p == nil {
		return &net.Dialer{}
	}
	d := &net.Dialer{
		Control: makeControl(p),
	}
	resolverDialer := func(ctx context.Context, network, address string) (net.Conn, error) {
		resolvers := strings.Split(p.GetResolvers(), ",")
		newAddress, err := replaceIP(address, resolvers)
		if err != nil {
			return nil, err
		}
		return d.DialContext(ctx, network, newAddress)
	}
	d.Resolver = &net.Resolver{
		PreferGo: true,
		Dial:     resolverDialer,
	}
	return d
}

// MakeListenConfig returns a new ListenConfig that creates protected
// listener sockets.
func MakeListenConfig(p Protector) *net.ListenConfig {
	if p == nil {
		return &net.ListenConfig{}
	}
	return &net.ListenConfig{
		Control: makeControl(p),
	}
}
<|endoftext|>