<commit_before>/*
Copyright IBM Corp. 2016 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 db

import (
	"bytes"
	"io/ioutil"
	"os"
	"testing"

	"github.com/spf13/viper"
	"github.com/tecbot/gorocksdb"
)

func TestMain(m *testing.M) {
	setupTestConfig()
	os.Exit(m.Run())
}

func TestCreateDB_DirDoesNotExist(t *testing.T) {
	err := CreateDB()
	if err != nil {
		t.Fatalf("Failed to create DB: %s", err)
	}
	deleteTestDB()
}

func TestCreateDB_NonEmptyDirExists(t *testing.T) {
	createNonEmptyTestDBPath()
	err := CreateDB()
	if err == nil {
		t.Fatal("Dir alrady exists. DB creation should throw error")
	}
	deleteTestDBPath()
}

func TestWriteAndRead(t *testing.T) {
	createTestDB()
	defer deleteTestDB()
	performBasicReadWrite(t)
}

func TestOpenDB_DirDoesNotExist(t *testing.T) {
	deleteTestDBPath()
	defer deleteTestDB()
	performBasicReadWrite(t)
}

func TestOpenDB_DirEmpty(t *testing.T) {
	deleteTestDBPath()
	createTestDBPath()
	defer deleteTestDB()
	performBasicReadWrite(t)
}

// This test verifies that when a new column family is added to the DB
// users at an older level of the DB will still be able to open it with new code
func TestDBColumnUpgrade(t *testing.T) {
	deleteTestDBPath()
	createTestDBPath()
	err := CreateDB()
	if nil != err {
		t.Fatalf("Error creating DB")
	}
	db, err := openDB()
	if nil != err {
		t.Fatalf("Error opening DB")
	}
	db.CloseDB()

	oldcfs := columnfamilies
	columnfamilies = append([]string{"Testing"}, columnfamilies...)
	defer func() {
		columnfamilies = oldcfs
	}()
	db, err = openDB()
	if nil != err {
		t.Fatalf("Error re-opening DB with upgraded columnFamilies")
	}
	db.CloseDB()
}

// db helper functions
func createTestDBPath() {
	dbPath := viper.GetString("peer.fileSystemPath")
	os.MkdirAll(dbPath, 0775)
}

func createNonEmptyTestDBPath() {
	dbPath := viper.GetString("peer.fileSystemPath")
	os.MkdirAll(dbPath+"/db/tmpFile", 0775)
}

func createTestDB() error {
	return CreateDB()
}

func deleteTestDBPath() {
	dbPath := viper.GetString("peer.fileSystemPath")
	os.RemoveAll(dbPath)
}

func deleteTestDB() {
	GetDBHandle().CloseDB()
	deleteTestDBPath()
}

func setupTestConfig() {
	tempDir, err := ioutil.TempDir("", "fabric-db-test")
	if err != nil {
		panic(err)
	}
	viper.Set("peer.fileSystemPath", tempDir)
	deleteTestDBPath()
}

func performBasicReadWrite(t *testing.T) {
	openchainDB := GetDBHandle()
	opt := gorocksdb.NewDefaultWriteOptions()
	defer opt.Destroy()
	writeBatch := gorocksdb.NewWriteBatch()
	defer writeBatch.Destroy()
	writeBatch.PutCF(openchainDB.BlockchainCF, []byte("dummyKey"), []byte("dummyValue"))
	err := openchainDB.DB.Write(opt, writeBatch)
	if err != nil {
		t.Fatal("Error while writing to db")
	}
	value, err := openchainDB.GetFromBlockchainCF([]byte("dummyKey"))

	if err != nil {
		t.Fatalf("read error = [%s]", err)
	}

	if !bytes.Equal(value, []byte("dummyValue")) {
		t.Fatal("read error. Bytes not equal")
	}
}
<commit_msg>added tests for enhancing code coverage in 'db' package<commit_after>/*
Copyright IBM Corp. 2016 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 db

import (
	"bytes"
	"io/ioutil"
	"os"
	"testing"

	"github.com/spf13/viper"
	"github.com/tecbot/gorocksdb"
)

func TestMain(m *testing.M) {
	setupTestConfig()
	os.Exit(m.Run())
}

func TestGetDBPathEmptyPath(t *testing.T) {
	originalSetting := viper.GetString("peer.fileSystemPath")
	viper.Set("peer.fileSystemPath", "")
	defer func() {
		x := recover()
		if x == nil {
			t.Fatal("A panic should have been caused here.")
		}
	}()
	defer viper.Set("peer.fileSystemPath", originalSetting)
	GetDBHandle()
}

func TestCreateDB_DirDoesNotExist(t *testing.T) {
	err := CreateDB()
	if err != nil {
		t.Fatalf("Failed to create DB: %s", err)
	}
	deleteTestDB()
}

func TestCreateDB_NonEmptyDirExists(t *testing.T) {
	createNonEmptyTestDBPath()
	err := CreateDB()
	if err == nil {
		t.Fatal("Dir alrady exists. DB creation should throw error")
	}
	deleteTestDBPath()
}

func TestWriteAndRead(t *testing.T) {
	createTestDB()
	defer deleteTestDB()
	performBasicReadWrite(t)
}

func TestOpenDB_DirDoesNotExist(t *testing.T) {
	deleteTestDBPath()
	defer deleteTestDB()
	performBasicReadWrite(t)
}

func TestOpenDB_DirEmpty(t *testing.T) {
	deleteTestDBPath()
	createTestDBPath()
	defer deleteTestDB()
	performBasicReadWrite(t)
}

// This test verifies that when a new column family is added to the DB
// users at an older level of the DB will still be able to open it with new code
func TestDBColumnUpgrade(t *testing.T) {
	deleteTestDBPath()
	createTestDBPath()
	err := CreateDB()
	if nil != err {
		t.Fatalf("Error creating DB")
	}
	db, err := openDB()
	if nil != err {
		t.Fatalf("Error opening DB")
	}
	db.CloseDB()

	oldcfs := columnfamilies
	columnfamilies = append([]string{"Testing"}, columnfamilies...)
	defer func() {
		columnfamilies = oldcfs
	}()
	db, err = openDB()
	if nil != err {
		t.Fatalf("Error re-opening DB with upgraded columnFamilies")
	}
	db.CloseDB()
}

func TestDeleteState(t *testing.T) {
	testDBWrapper := NewTestDBWrapper()
	testDBWrapper.CreateFreshDB(t)
	openchainDB := GetDBHandle()
	defer testDBWrapper.cleanup()
	openchainDB.Put(openchainDB.StateCF, []byte("key1"), []byte("value1"))
	openchainDB.Put(openchainDB.StateDeltaCF, []byte("key2"), []byte("value2"))
	openchainDB.DeleteState()
	value1, err := openchainDB.GetFromStateCF([]byte("key1"))
	if err != nil {
		t.Fatalf("Error getting in value: %s", err)
	}
	if value1 != nil {
		t.Fatalf("A nil value expected. Found [%s]", value1)
	}

	value2, err := openchainDB.GetFromStateCF([]byte("key2"))
	if err != nil {
		t.Fatalf("Error getting in value: %s", err)
	}
	if value2 != nil {
		t.Fatalf("A nil value expected. Found [%s]", value2)
	}
}

func TestDBSnapshot(t *testing.T) {
	testDBWrapper := NewTestDBWrapper()
	testDBWrapper.CreateFreshDB(t)
	openchainDB := GetDBHandle()
	defer testDBWrapper.cleanup()

	// write key-values
	openchainDB.Put(openchainDB.BlockchainCF, []byte("key1"), []byte("value1"))
	openchainDB.Put(openchainDB.BlockchainCF, []byte("key2"), []byte("value2"))

	// create a snapshot
	snapshot := openchainDB.GetSnapshot()

	// add/delete/modify key-values
	openchainDB.Delete(openchainDB.BlockchainCF, []byte("key1"))
	openchainDB.Put(openchainDB.BlockchainCF, []byte("key2"), []byte("value2_new"))
	openchainDB.Put(openchainDB.BlockchainCF, []byte("key3"), []byte("value3"))

	// test key-values from latest data in db
	v1, _ := openchainDB.GetFromBlockchainCF([]byte("key1"))
	v2, _ := openchainDB.GetFromBlockchainCF([]byte("key2"))
	v3, _ := openchainDB.GetFromBlockchainCF([]byte("key3"))
	if !bytes.Equal(v1, nil) {
		t.Fatalf("Expected value from db is 'nil', found [%s]", v1)
	}
	if !bytes.Equal(v2, []byte("value2_new")) {
		t.Fatalf("Expected value from db [%s], found [%s]", "value2_new", v2)
	}
	if !bytes.Equal(v3, []byte("value3")) {
		t.Fatalf("Expected value from db [%s], found [%s]", "value3", v3)
	}

	// test key-values from snapshot
	v1, _ = openchainDB.GetFromBlockchainCFSnapshot(snapshot, []byte("key1"))
	v2, _ = openchainDB.GetFromBlockchainCFSnapshot(snapshot, []byte("key2"))
	v3, err := openchainDB.GetFromBlockchainCFSnapshot(snapshot, []byte("key3"))
	if err != nil {
		t.Fatalf("Error: %s", err)
	}

	if !bytes.Equal(v1, []byte("value1")) {
		t.Fatalf("Expected value from db snapshot [%s], found [%s]", "value1", v1)
	}

	if !bytes.Equal(v2, []byte("value2")) {
		t.Fatalf("Expected value from db snapshot [%s], found [%s]", "value1", v2)
	}

	if !bytes.Equal(v3, nil) {
		t.Fatalf("Expected value from db snapshot is 'nil', found [%s]", v3)
	}
}

func TestDBIteratorAndSnapshotIterator(t *testing.T) {
	testDBWrapper := NewTestDBWrapper()
	testDBWrapper.CreateFreshDB(t)
	openchainDB := GetDBHandle()
	defer testDBWrapper.cleanup()

	// write key-values
	openchainDB.Put(openchainDB.StateCF, []byte("key1"), []byte("value1"))
	openchainDB.Put(openchainDB.StateCF, []byte("key2"), []byte("value2"))

	// create a snapshot
	snapshot := openchainDB.GetSnapshot()

	// add/delete/modify key-values
	openchainDB.Delete(openchainDB.StateCF, []byte("key1"))
	openchainDB.Put(openchainDB.StateCF, []byte("key2"), []byte("value2_new"))
	openchainDB.Put(openchainDB.StateCF, []byte("key3"), []byte("value3"))

	// test snapshot iterator
	itr := openchainDB.GetStateCFSnapshotIterator(snapshot)
	defer itr.Close()
	testIterator(t, itr, map[string][]byte{"key1": []byte("value1"), "key2": []byte("value2")})

	// test iterator over latest data in stateCF
	itr = openchainDB.GetStateCFIterator()
	defer itr.Close()
	testIterator(t, itr, map[string][]byte{"key2": []byte("value2_new"), "key3": []byte("value3")})

	openchainDB.Put(openchainDB.StateDeltaCF, []byte("key4"), []byte("value4"))
	openchainDB.Put(openchainDB.StateDeltaCF, []byte("key5"), []byte("value5"))
	itr = openchainDB.GetStateDeltaCFIterator()
	defer itr.Close()
	testIterator(t, itr, map[string][]byte{"key4": []byte("value4"), "key5": []byte("value5")})

	openchainDB.Put(openchainDB.BlockchainCF, []byte("key6"), []byte("value6"))
	openchainDB.Put(openchainDB.BlockchainCF, []byte("key7"), []byte("value7"))
	itr = openchainDB.GetBlockchainCFIterator()
	defer itr.Close()
	testIterator(t, itr, map[string][]byte{"key6": []byte("value6"), "key7": []byte("value7")})
}

func testIterator(t *testing.T, itr *gorocksdb.Iterator, expectedValues map[string][]byte) {
	itrResults := make(map[string][]byte)
	itr.SeekToFirst()
	for ; itr.Valid(); itr.Next() {
		key := itr.Key()
		value := itr.Value()
		k := makeCopy(key.Data())
		v := makeCopy(value.Data())
		itrResults[string(k)] = v
	}
	if len(itrResults) != len(expectedValues) {
		t.Fatalf("Expected [%d] results from iterator, found [%d]", len(expectedValues), len(itrResults))
	}
	for k, v := range expectedValues {
		if !bytes.Equal(itrResults[k], v) {
			t.Fatalf("Wrong value for key [%s]. Expected [%s], found [%s]", k, itrResults[k], v)
		}
	}
}

// db helper functions
func createTestDBPath() {
	dbPath := viper.GetString("peer.fileSystemPath")
	os.MkdirAll(dbPath, 0775)
}

func createNonEmptyTestDBPath() {
	dbPath := viper.GetString("peer.fileSystemPath")
	os.MkdirAll(dbPath+"/db/tmpFile", 0775)
}

func createTestDB() error {
	return CreateDB()
}

func deleteTestDBPath() {
	dbPath := viper.GetString("peer.fileSystemPath")
	os.RemoveAll(dbPath)
}

func deleteTestDB() {
	GetDBHandle().CloseDB()
	deleteTestDBPath()
}

func setupTestConfig() {
	tempDir, err := ioutil.TempDir("", "fabric-db-test")
	if err != nil {
		panic(err)
	}
	viper.Set("peer.fileSystemPath", tempDir)
	deleteTestDBPath()
}

func performBasicReadWrite(t *testing.T) {
	openchainDB := GetDBHandle()
	opt := gorocksdb.NewDefaultWriteOptions()
	defer opt.Destroy()
	writeBatch := gorocksdb.NewWriteBatch()
	defer writeBatch.Destroy()
	writeBatch.PutCF(openchainDB.BlockchainCF, []byte("dummyKey"), []byte("dummyValue"))
	writeBatch.PutCF(openchainDB.StateCF, []byte("dummyKey1"), []byte("dummyValue1"))
	writeBatch.PutCF(openchainDB.StateDeltaCF, []byte("dummyKey2"), []byte("dummyValue2"))
	writeBatch.PutCF(openchainDB.IndexesCF, []byte("dummyKey3"), []byte("dummyValue3"))
	err := openchainDB.DB.Write(opt, writeBatch)
	if err != nil {
		t.Fatalf("Error while writing to db: %s", err)
	}
	value, err := openchainDB.GetFromBlockchainCF([]byte("dummyKey"))
	if err != nil {
		t.Fatalf("read error = [%s]", err)
	}
	if !bytes.Equal(value, []byte("dummyValue")) {
		t.Fatalf("read error. Bytes not equal. Expected [%s], found [%s]", "dummyValue", value)
	}

	value, err = openchainDB.GetFromStateCF([]byte("dummyKey1"))
	if err != nil {
		t.Fatalf("read error = [%s]", err)
	}
	if !bytes.Equal(value, []byte("dummyValue1")) {
		t.Fatalf("read error. Bytes not equal. Expected [%s], found [%s]", "dummyValue1", value)
	}

	value, err = openchainDB.GetFromStateDeltaCF([]byte("dummyKey2"))
	if err != nil {
		t.Fatalf("read error = [%s]", err)
	}
	if !bytes.Equal(value, []byte("dummyValue2")) {
		t.Fatalf("read error. Bytes not equal. Expected [%s], found [%s]", "dummyValue2", value)
	}

	value, err = openchainDB.GetFromIndexesCF([]byte("dummyKey3"))
	if err != nil {
		t.Fatalf("read error = [%s]", err)
	}
	if !bytes.Equal(value, []byte("dummyValue3")) {
		t.Fatalf("read error. Bytes not equal. Expected [%s], found [%s]", "dummyValue3", value)
	}
}
<|endoftext|>


<commit_before>package core

import (
	"fmt"
	"reflect"
	"testing"

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

func TestWrite_returnsErrorIfTargetNotPtr(t *testing.T) {
	// try to copy a value to a non-pointer
	err := WriteAnswer(true, "hello", true)
	// make sure there was an error
	if err == nil {
		t.Error("Did not encounter error when writing to non-pointer.")
	}
}

func TestWrite_canWriteToBool(t *testing.T) {
	// a pointer to hold the boolean value
	ptr := true

	// try to copy a false value to the pointer
	WriteAnswer(&ptr, "", false)

	// if the value is true
	if ptr {
		// the test failed
		t.Error("Could not write a false bool to a pointer")
	}
}

func TestWrite_canWriteString(t *testing.T) {
	// a pointer to hold the boolean value
	ptr := ""

	// try to copy a false value to the pointer
	err := WriteAnswer(&ptr, "", "hello")
	if err != nil {
		t.Error(err)
	}

	// if the value is not what we wrote
	if ptr != "hello" {
		t.Error("Could not write a string value to a pointer")
	}
}

func TestWrite_canWriteSlice(t *testing.T) {
	// a pointer to hold the value
	ptr := []string{}

	// copy in a value
	WriteAnswer(&ptr, "", []string{"hello", "world"})

	// make sure there are two entries
	if len(ptr) != 2 {
		// the test failed
		t.Errorf("Incorrect number of entries in written list. Expected 2, found %v.", len(ptr))
		// dont move on
		return
	}

	// make sure the first entry is hello
	if ptr[0] != "hello" {
		// the test failed
		t.Errorf("incorrect first value in written pointer. expected hello found %v.", ptr[0])
	}

	// make sure the second entry is world
	if ptr[1] != "world" {
		// the test failed
		t.Errorf("incorrect second value in written pointer. expected world found %v.", ptr[0])
	}
}

func TestWrite_recoversInvalidReflection(t *testing.T) {
	// a variable to mutate
	ptr := false

	// write a boolean value to the string
	err := WriteAnswer(&ptr, "", "hello")

	// if there was no error
	if err == nil {
		// the test failed
		t.Error("Did not encounter error when forced invalid write.")
	}
}

func TestWriteAnswer_handlesNonStructValues(t *testing.T) {
	// the value to write to
	ptr := ""

	// write a value to the pointer
	WriteAnswer(&ptr, "", "world")

	// if we didn't change the value appropriate
	if ptr != "world" {
		// the test failed
		t.Error("Did not write value to primitive pointer")
	}
}

func TestWriteAnswer_canMutateStruct(t *testing.T) {
	// the struct to hold the answer
	ptr := struct{ Name string }{}

	// write a value to an existing field
	err := WriteAnswer(&ptr, "name", "world")
	if err != nil {
		// the test failed
		t.Errorf("Encountered error while writing answer: %v", err.Error())
		// we're done here
		return
	}

	// make sure we changed the field
	if ptr.Name != "world" {
		// the test failed
		t.Error("Did not mutate struct field when writing answer.")
	}
}

func TestWriteAnswer_canMutateMap(t *testing.T) {
	// the map to hold the answer
	ptr := make(map[string]interface{})

	// write a value to an existing field
	err := WriteAnswer(&ptr, "name", "world")
	if err != nil {
		// the test failed
		t.Errorf("Encountered error while writing answer: %v", err.Error())
		// we're done here
		return
	}

	// make sure we changed the field
	if ptr["name"] != "world" {
		// the test failed
		t.Error("Did not mutate map when writing answer.")
	}
}

func TestWrite_returnsErrorIfInvalidMapType(t *testing.T) {
	// try to copy a value to a non map[string]interface{}
	ptr := make(map[int]string)

	err := WriteAnswer(&ptr, "name", "world")
	// make sure there was an error
	if err == nil {
		t.Error("Did not encounter error when writing to invalid map.")
	}
}

func TestWriteAnswer_returnsErrWhenFieldNotFound(t *testing.T) {
	// the struct to hold the answer
	ptr := struct{ Name string }{}

	// write a value to an existing field
	err := WriteAnswer(&ptr, "", "world")

	if err == nil {
		// the test failed
		t.Error("Did not encountered error while writing answer to non-existing field.")
	}
}

func TestFindFieldIndex_canFindExportedField(t *testing.T) {
	// create a reflective wrapper over the struct to look through
	val := reflect.ValueOf(struct{ Name string }{})

	// find the field matching "name"
	fieldIndex, err := findFieldIndex(val, "name")
	// if something went wrong
	if err != nil {
		// the test failed
		t.Error(err.Error())
		return
	}

	// make sure we got the right value
	if val.Type().Field(fieldIndex).Name != "Name" {
		// the test failed
		t.Errorf("Did not find the correct field name. Expected 'Name' found %v.", val.Type().Field(fieldIndex).Name)
	}
}

func TestFindFieldIndex_canFindTaggedField(t *testing.T) {
	// the struct to look through
	val := reflect.ValueOf(struct {
		Username string `survey:"name"`
	}{})

	// find the field matching "name"
	fieldIndex, err := findFieldIndex(val, "name")
	// if something went wrong
	if err != nil {
		// the test failed
		t.Error(err.Error())
		return
	}

	// make sure we got the right value
	if val.Type().Field(fieldIndex).Name != "Username" {
		// the test failed
		t.Errorf("Did not find the correct field name. Expected 'Username' found %v.", val.Type().Field(fieldIndex).Name)
	}
}

func TestFindFieldIndex_canHandleCapitalAnswerNames(t *testing.T) {
	// create a reflective wrapper over the struct to look through
	val := reflect.ValueOf(struct{ Name string }{})

	// find the field matching "name"
	fieldIndex, err := findFieldIndex(val, "Name")
	// if something went wrong
	if err != nil {
		// the test failed
		t.Error(err.Error())
		return
	}

	// make sure we got the right value
	if val.Type().Field(fieldIndex).Name != "Name" {
		// the test failed
		t.Errorf("Did not find the correct field name. Expected 'Name' found %v.", val.Type().Field(fieldIndex).Name)
	}
}

func TestFindFieldIndex_tagOverwriteFieldName(t *testing.T) {
	// the struct to look through
	val := reflect.ValueOf(struct {
		Name     string
		Username string `survey:"name"`
	}{})

	// find the field matching "name"
	fieldIndex, err := findFieldIndex(val, "name")
	// if something went wrong
	if err != nil {
		// the test failed
		t.Error(err.Error())
		return
	}

	// make sure we got the right value
	if val.Type().Field(fieldIndex).Name != "Username" {
		// the test failed
		t.Errorf("Did not find the correct field name. Expected 'Username' found %v.", val.Type().Field(fieldIndex).Name)
	}
}

type testFieldSettable struct {
	Values map[string]string
}

func (t *testFieldSettable) WriteAnswer(name string, value interface{}) error {
	if t.Values == nil {
		t.Values = map[string]string{}
	}
	if v, ok := value.(string); ok {
		t.Values[name] = v
		return nil
	}
	return fmt.Errorf("Incompatible type %T", value)
}

func TestWriteWithFieldSettable(t *testing.T) {
	testSet1 := testFieldSettable{}
	err := WriteAnswer(&testSet1, "values", "stringVal")
	assert.Nil(t, err)
	assert.Equal(t, map[string]string{"values": "stringVal"}, testSet1.Values)

	testSet2 := testFieldSettable{}
	err = WriteAnswer(&testSet2, "values", 123)
	assert.Error(t, fmt.Errorf("Incompatible type int64"), err)
	assert.Equal(t, map[string]string{}, testSet2.Values)
}

// CONVERSION TESTS
func TestWrite_canStringToBool(t *testing.T) {
	// a pointer to hold the boolean value
	ptr := true

	// try to copy a false value to the pointer
	WriteAnswer(&ptr, "", "false")

	// if the value is true
	if ptr {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToInt(t *testing.T) {
	// a pointer to hold the value
	var ptr int = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToInt8(t *testing.T) {
	// a pointer to hold the value
	var ptr int8 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToInt16(t *testing.T) {
	// a pointer to hold the value
	var ptr int16 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToInt32(t *testing.T) {
	// a pointer to hold the value
	var ptr int32 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToInt64(t *testing.T) {
	// a pointer to hold the value
	var ptr int64 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToUint(t *testing.T) {
	// a pointer to hold the value
	var ptr uint = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToUint8(t *testing.T) {
	// a pointer to hold the value
	var ptr uint8 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToUint16(t *testing.T) {
	// a pointer to hold the value
	var ptr uint16 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToUint32(t *testing.T) {
	// a pointer to hold the value
	var ptr uint32 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToUint64(t *testing.T) {
	// a pointer to hold the value
	var ptr uint64 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToFloat32(t *testing.T) {
	// a pointer to hold the value
	var ptr float32 = 1.0

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2.5")

	// if the value is true
	if ptr != 2.5 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToFloat64(t *testing.T) {
	// a pointer to hold the value
	var ptr float64 = 1.0

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2.5")

	// if the value is true
	if ptr != 2.5 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canConvertStructFieldTypes(t *testing.T) {
	// the struct to hold the answer
	ptr := struct {
		Name   string
		Age    uint
		Male   bool
		Height float64
	}{}

	// write the values as strings
	check(t, WriteAnswer(&ptr, "name", "Bob"))
	check(t, WriteAnswer(&ptr, "age", "22"))
	check(t, WriteAnswer(&ptr, "male", "true"))
	check(t, WriteAnswer(&ptr, "height", "6.2"))

	// make sure we changed the fields
	if ptr.Name != "Bob" {
		t.Error("Did not mutate Name when writing answer.")
	}

	if ptr.Age != 22 {
		t.Error("Did not mutate Age when writing answer.")
	}

	if !ptr.Male {
		t.Error("Did not mutate Male when writing answer.")
	}

	if ptr.Height != 6.2 {
		t.Error("Did not mutate Height when writing answer.")
	}
}

func check(t *testing.T, err error) {
	if err != nil {
		t.Fatalf("Encountered error while writing answer: %v", err.Error())
	}
}
<commit_msg>added test for writing to tagged fields<commit_after>package core

import (
	"fmt"
	"reflect"
	"testing"

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

func TestWrite_returnsErrorIfTargetNotPtr(t *testing.T) {
	// try to copy a value to a non-pointer
	err := WriteAnswer(true, "hello", true)
	// make sure there was an error
	if err == nil {
		t.Error("Did not encounter error when writing to non-pointer.")
	}
}

func TestWrite_canWriteToBool(t *testing.T) {
	// a pointer to hold the boolean value
	ptr := true

	// try to copy a false value to the pointer
	WriteAnswer(&ptr, "", false)

	// if the value is true
	if ptr {
		// the test failed
		t.Error("Could not write a false bool to a pointer")
	}
}

func TestWrite_canWriteString(t *testing.T) {
	// a pointer to hold the boolean value
	ptr := ""

	// try to copy a false value to the pointer
	err := WriteAnswer(&ptr, "", "hello")
	if err != nil {
		t.Error(err)
	}

	// if the value is not what we wrote
	if ptr != "hello" {
		t.Error("Could not write a string value to a pointer")
	}
}

func TestWrite_canWriteSlice(t *testing.T) {
	// a pointer to hold the value
	ptr := []string{}

	// copy in a value
	WriteAnswer(&ptr, "", []string{"hello", "world"})

	// make sure there are two entries
	if len(ptr) != 2 {
		// the test failed
		t.Errorf("Incorrect number of entries in written list. Expected 2, found %v.", len(ptr))
		// dont move on
		return
	}

	// make sure the first entry is hello
	if ptr[0] != "hello" {
		// the test failed
		t.Errorf("incorrect first value in written pointer. expected hello found %v.", ptr[0])
	}

	// make sure the second entry is world
	if ptr[1] != "world" {
		// the test failed
		t.Errorf("incorrect second value in written pointer. expected world found %v.", ptr[0])
	}
}

func TestWrite_recoversInvalidReflection(t *testing.T) {
	// a variable to mutate
	ptr := false

	// write a boolean value to the string
	err := WriteAnswer(&ptr, "", "hello")

	// if there was no error
	if err == nil {
		// the test failed
		t.Error("Did not encounter error when forced invalid write.")
	}
}

func TestWriteAnswer_handlesNonStructValues(t *testing.T) {
	// the value to write to
	ptr := ""

	// write a value to the pointer
	WriteAnswer(&ptr, "", "world")

	// if we didn't change the value appropriate
	if ptr != "world" {
		// the test failed
		t.Error("Did not write value to primitive pointer")
	}
}

func TestWriteAnswer_canMutateStruct(t *testing.T) {
	// the struct to hold the answer
	ptr := struct{ Name string }{}

	// write a value to an existing field
	err := WriteAnswer(&ptr, "name", "world")
	if err != nil {
		// the test failed
		t.Errorf("Encountered error while writing answer: %v", err.Error())
		// we're done here
		return
	}

	// make sure we changed the field
	if ptr.Name != "world" {
		// the test failed
		t.Error("Did not mutate struct field when writing answer.")
	}
}

func TestWriteAnswer_canMutateMap(t *testing.T) {
	// the map to hold the answer
	ptr := make(map[string]interface{})

	// write a value to an existing field
	err := WriteAnswer(&ptr, "name", "world")
	if err != nil {
		// the test failed
		t.Errorf("Encountered error while writing answer: %v", err.Error())
		// we're done here
		return
	}

	// make sure we changed the field
	if ptr["name"] != "world" {
		// the test failed
		t.Error("Did not mutate map when writing answer.")
	}
}

func TestWrite_returnsErrorIfInvalidMapType(t *testing.T) {
	// try to copy a value to a non map[string]interface{}
	ptr := make(map[int]string)

	err := WriteAnswer(&ptr, "name", "world")
	// make sure there was an error
	if err == nil {
		t.Error("Did not encounter error when writing to invalid map.")
	}
}

func TestWriteAnswer_returnsErrWhenFieldNotFound(t *testing.T) {
	// the struct to hold the answer
	ptr := struct{ Name string }{}

	// write a value to an existing field
	err := WriteAnswer(&ptr, "", "world")

	if err == nil {
		// the test failed
		t.Error("Did not encountered error while writing answer to non-existing field.")
	}
}

func TestFindFieldIndex_canFindExportedField(t *testing.T) {
	// create a reflective wrapper over the struct to look through
	val := reflect.ValueOf(struct{ Name string }{})

	// find the field matching "name"
	fieldIndex, err := findFieldIndex(val, "name")
	// if something went wrong
	if err != nil {
		// the test failed
		t.Error(err.Error())
		return
	}

	// make sure we got the right value
	if val.Type().Field(fieldIndex).Name != "Name" {
		// the test failed
		t.Errorf("Did not find the correct field name. Expected 'Name' found %v.", val.Type().Field(fieldIndex).Name)
	}
}

func TestFindFieldIndex_canFindTaggedField(t *testing.T) {
	// the struct to look through
	val := reflect.ValueOf(struct {
		Username string `survey:"name"`
	}{})

	// find the field matching "name"
	fieldIndex, err := findFieldIndex(val, "name")
	// if something went wrong
	if err != nil {
		// the test failed
		t.Error(err.Error())
		return
	}

	// make sure we got the right value
	if val.Type().Field(fieldIndex).Name != "Username" {
		// the test failed
		t.Errorf("Did not find the correct field name. Expected 'Username' found %v.", val.Type().Field(fieldIndex).Name)
	}
}

func TestFindFieldIndex_canHandleCapitalAnswerNames(t *testing.T) {
	// create a reflective wrapper over the struct to look through
	val := reflect.ValueOf(struct{ Name string }{})

	// find the field matching "name"
	fieldIndex, err := findFieldIndex(val, "Name")
	// if something went wrong
	if err != nil {
		// the test failed
		t.Error(err.Error())
		return
	}

	// make sure we got the right value
	if val.Type().Field(fieldIndex).Name != "Name" {
		// the test failed
		t.Errorf("Did not find the correct field name. Expected 'Name' found %v.", val.Type().Field(fieldIndex).Name)
	}
}

func TestFindFieldIndex_tagOverwriteFieldName(t *testing.T) {
	// the struct to look through
	val := reflect.ValueOf(struct {
		Name     string
		Username string `survey:"name"`
	}{})

	// find the field matching "name"
	fieldIndex, err := findFieldIndex(val, "name")
	// if something went wrong
	if err != nil {
		// the test failed
		t.Error(err.Error())
		return
	}

	// make sure we got the right value
	if val.Type().Field(fieldIndex).Name != "Username" {
		// the test failed
		t.Errorf("Did not find the correct field name. Expected 'Username' found %v.", val.Type().Field(fieldIndex).Name)
	}
}

type testFieldSettable struct {
	Values map[string]string
}

type testTaggedStruct struct {
	TaggedValue string `survey:"tagged"`
}

func (t *testFieldSettable) WriteAnswer(name string, value interface{}) error {
	if t.Values == nil {
		t.Values = map[string]string{}
	}
	if v, ok := value.(string); ok {
		t.Values[name] = v
		return nil
	}
	return fmt.Errorf("Incompatible type %T", value)
}

func TestWriteWithFieldSettable(t *testing.T) {
	testSet1 := testFieldSettable{}
	err := WriteAnswer(&testSet1, "values", "stringVal")
	assert.Nil(t, err)
	assert.Equal(t, map[string]string{"values": "stringVal"}, testSet1.Values)

	testSet2 := testFieldSettable{}
	err = WriteAnswer(&testSet2, "values", 123)
	assert.Error(t, fmt.Errorf("Incompatible type int64"), err)
	assert.Equal(t, map[string]string{}, testSet2.Values)

	testSetStruct := testTaggedStruct{}
	err = WriteAnswer(&testSetStruct, "tagged", "stringVal1")
	assert.Nil(t, err)
	assert.Equal(t, testSetStruct.TaggedValue, "stringVal1")
}

// CONVERSION TESTS
func TestWrite_canStringToBool(t *testing.T) {
	// a pointer to hold the boolean value
	ptr := true

	// try to copy a false value to the pointer
	WriteAnswer(&ptr, "", "false")

	// if the value is true
	if ptr {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToInt(t *testing.T) {
	// a pointer to hold the value
	var ptr int = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToInt8(t *testing.T) {
	// a pointer to hold the value
	var ptr int8 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToInt16(t *testing.T) {
	// a pointer to hold the value
	var ptr int16 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToInt32(t *testing.T) {
	// a pointer to hold the value
	var ptr int32 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToInt64(t *testing.T) {
	// a pointer to hold the value
	var ptr int64 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToUint(t *testing.T) {
	// a pointer to hold the value
	var ptr uint = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToUint8(t *testing.T) {
	// a pointer to hold the value
	var ptr uint8 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToUint16(t *testing.T) {
	// a pointer to hold the value
	var ptr uint16 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToUint32(t *testing.T) {
	// a pointer to hold the value
	var ptr uint32 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToUint64(t *testing.T) {
	// a pointer to hold the value
	var ptr uint64 = 1

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2")

	// if the value is true
	if ptr != 2 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToFloat32(t *testing.T) {
	// a pointer to hold the value
	var ptr float32 = 1.0

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2.5")

	// if the value is true
	if ptr != 2.5 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canStringToFloat64(t *testing.T) {
	// a pointer to hold the value
	var ptr float64 = 1.0

	// try to copy a value to the pointer
	WriteAnswer(&ptr, "", "2.5")

	// if the value is true
	if ptr != 2.5 {
		// the test failed
		t.Error("Could not convert string to pointer type")
	}
}

func TestWrite_canConvertStructFieldTypes(t *testing.T) {
	// the struct to hold the answer
	ptr := struct {
		Name   string
		Age    uint
		Male   bool
		Height float64
	}{}

	// write the values as strings
	check(t, WriteAnswer(&ptr, "name", "Bob"))
	check(t, WriteAnswer(&ptr, "age", "22"))
	check(t, WriteAnswer(&ptr, "male", "true"))
	check(t, WriteAnswer(&ptr, "height", "6.2"))

	// make sure we changed the fields
	if ptr.Name != "Bob" {
		t.Error("Did not mutate Name when writing answer.")
	}

	if ptr.Age != 22 {
		t.Error("Did not mutate Age when writing answer.")
	}

	if !ptr.Male {
		t.Error("Did not mutate Male when writing answer.")
	}

	if ptr.Height != 6.2 {
		t.Error("Did not mutate Height when writing answer.")
	}
}

func check(t *testing.T, err error) {
	if err != nil {
		t.Fatalf("Encountered error while writing answer: %v", err.Error())
	}
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net"
	"time"
	"flag"
	"path"
	"strings"
	"os"
	"os/exec"
	"os/signal"
	"syscall"

	"github.com/miekg/dns"
)

var nameservers []string
var verbose *bool
var gateway4 *string
var gateway6 *string
var routev6 bool
var router string
var routedv4 map[string]struct{}
var routedv6 map[string]struct{}

func runAndLog(name string, arg ...string) {
	out, err := exec.Command(name, arg...).CombinedOutput()

	if err != nil {
		if len(out) > 1 {
			log.Printf("%s: %s", path.Base(strings.Replace(name, "\\", "/", -1)), out)
		} else {
			log.Print(err)
		}
	} else if *verbose {
		log.Printf("%s: %s", path.Base(strings.Replace(name, "\\", "/", -1)), out)
	}
}

func getEmptyMsg(w dns.ResponseWriter, req *dns.Msg, err int) *dns.Msg {
	m := new(dns.Msg)

	m.SetReply(req)

	if err != 0 {
		m.SetRcode(req, err)
	}

	m.Authoritative = false
	m.RecursionAvailable = true

	return m
}

func getServerReply(w dns.ResponseWriter, req *dns.Msg) *dns.Msg {
	if *verbose {
		log.Print("Forwarding ", req.Question[0].Name, "/", dns.Type(req.Question[0].Qtype).String())
	}

	client := &dns.Client{Net: "udp", ReadTimeout: 4 * time.Second, WriteTimeout: 4 * time.Second, SingleInflight: true}

	if _, tcp := w.RemoteAddr().(*net.TCPAddr); tcp {
		client.Net = "tcp"
	}

	var r *dns.Msg
	var err error

	for i := 0; i < len(nameservers); i++ {
		r, _, err = client.Exchange(req, nameservers[(int(req.Id) + i) % len(nameservers)])
		if err == nil {
			r.Compress = true
			return r
		}
	}

	log.Print("Failed to forward request.", err)
	return getEmptyMsg(w, req, dns.RcodeServerFailure)
}

func handleRequest(w dns.ResponseWriter, req *dns.Msg) {
	var m *dns.Msg

	if len(req.Question) > 0 && (req.Question[0].Name == "netflix.com." || strings.HasSuffix(req.Question[0].Name, ".netflix.com.")) {
		if req.Question[0].Qtype == dns.TypeA {
			m = getServerReply(w, req)
			for _, ans := range m.Answer {
				if ans.Header().Rrtype == dns.TypeA {
					ip := ans.(*dns.A).A.String()
					routedv4[ip] = struct{}{}

					log.Print("Re-routing ", ip, " for ", ans.Header().Name, "/", dns.Type(ans.Header().Rrtype).String())

					runAndLog(router, "add", ip + "/32", *gateway4)
				} else if ans.Header().Rrtype == dns.TypeAAAA {
					// sanity check for now, shouldn't happen afaik
					log.Print("WARNING: AAAA response in ", ans.Header().Name, "/A")
				}
			}
		} else if req.Question[0].Qtype == dns.TypeAAAA {
			if routev6 {
				m = getServerReply(w, req)
				for _, ans := range m.Answer {
					if ans.Header().Rrtype == dns.TypeAAAA {
						ip := ans.(*dns.AAAA).AAAA.String()
						routedv6[ip] = struct{}{}

						log.Print("Re-routing ", ip, " for ", ans.Header().Name, "/", dns.Type(ans.Header().Rrtype).String())

						runAndLog(router, "add", ip + "/128", *gateway6)
					} else if ans.Header().Rrtype == dns.TypeA {
						log.Print("WARNING: A response in ", ans.Header().Name, "/AAAA")
					}
				}
			} else {
				if *verbose {
					log.Print("Hijacking ", req.Question[0].Name, "/", dns.Type(req.Question[0].Qtype).String())
				}
				m = getEmptyMsg(w, req, 0)
			}
		} else {
			m = getServerReply(w, req)
		}
	} else {
		m = getServerReply(w, req)
	}

	w.WriteMsg(m)
}

func removeRoutes() {
	if len(routedv4) > 0 {
		log.Print("Removing routes...")

		for ip, _ := range routedv4 {
			runAndLog(router, "delete", ip + "/32")
		}
	}

	if routev6 && len(routedv6) > 0 {
		log.Print("Removing IPv6 routes...")

		for ip, _ := range routedv6 {
			runAndLog(router, "delete", ip + "/128")
		}
	}
}

func main() {
	gateway4 = flag.String("r", "", "IPv4 gateway address for routing destination")
	gateway6 = flag.String("r6", "", "IPv6 gateway address for routing destination")
	verbose = flag.Bool("v", false, "verbose logging")

	flag.Usage = func() {
		flag.PrintDefaults()
	}
	flag.Parse()

	router, _ = exec.LookPath("route")
	if len(router) < 1 {
		log.Fatal("Unable to find the `route` command in your %PATH%.")
	}

	if len(*gateway4) < 1 {
		log.Fatal("A gateway IP must be specified via argument `r`.")
	}

	if ip := net.ParseIP(*gateway4); ip != nil && ip.To4() != nil {
		*gateway4 = ip.String()
	} else {
		log.Fatal("Specified gateway IP is not a valid IPv4 address.")
	}

	if len(*gateway6) > 1 {
		if ip := net.ParseIP(*gateway6); ip != nil && ip.To4() == nil {
			*gateway6 = ip.String()
			  routev6 = true
		} else {
			log.Fatal("Specified gateway IP is not a valid IPv6 address.")
		}
	} else {
		routev6 = false
	}

	routedv4 = make(map[string]struct{})

	if (routev6) {
		routedv6 = make(map[string]struct{})
	}

	log.Print("Starting DNS resolver...")

	nameservers = []string{"8.8.8.8:53", "8.8.4.4:53"}

	log.Print("Forwarding to ", nameservers)

	dns.HandleFunc(".", handleRequest)

	go func() {
		srv := &dns.Server{Addr: ":53", Net: "udp"}
		err := srv.ListenAndServe()
		if err != nil {
			log.Fatal("Failed to start UDP server.", err.Error())
		}
	}()

	go func() {
		srv := &dns.Server{Addr: ":53", Net: "tcp"}
		err := srv.ListenAndServe()
		if err != nil {
			log.Fatal("Failed to start TCP server.", err.Error())
		}
	}()

	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)

	for {
		select {
		case s := <-sig:
			removeRoutes()
			log.Fatalf("Received signal %d, exiting...", s)
		}
	}
}<commit_msg>Added documentation to the Go file.<commit_after>package main

import (
	"log"
	"net"
	"time"
	"flag"
	"path"
	"strings"
	"os"
	"os/exec"
	"os/signal"
	"syscall"

	"github.com/miekg/dns"
)

var (
	nameservers []string
	verbose     *bool
	gateway4    *string
	gateway6    *string
	dnsr1       *string
	dnsr2       *string
	routev6     bool
	router      string
	routedv4    map[string]struct{}
	routedv6    map[string]struct{}
)

// Executes the specified command with the specified arguments. The
// `name` parameter should be an absolute path to the executable.
// See `exec.LookPath()` for resolving names found in `%PATH%`.
// Any errors will be logged to stdout, if an output is available,
// otherwise the return code or internal Go error will be displayed.
func runAndLog(name string, arg ...string) {
	out, err := exec.Command(name, arg...).CombinedOutput()

	if err != nil {
		if len(out) > 1 {
			log.Printf("%s: %s", path.Base(strings.Replace(name, "\\", "/", -1)), out)
		} else {
			log.Print(err)
		}
	} else if *verbose {
		log.Printf("%s: %s", path.Base(strings.Replace(name, "\\", "/", -1)), out)
	}
}

// Creates an empty response to the specified request. If `err` is
// specified, the `RCODE` field in the response will be set to this value.
// If `err` is set to 0, the `RCODE` field will not be modified, and the
// resulting packet will just mean that the domain exists (not `NXDOMAIN`)
// but there are no records of the requested type associated to it.
// If `NXDOMAIN` is sent as a reply to the hijacked `AAAA` records of hostnames
// when IPv6 routing is disabled, some web browsers (e.g. Chrome) will display
// an error message stating `DNS_PROBE_FINISHED_NXDOMAIN`, even though a request
// for `A` records types is sent and properly replied to by the server to the client.
// Even though the original `ResponseWriter` object is taken as an argument,
// this function does not send a reply to the client. Instead, the
// packet is returned for further processing by the caller.
func getEmptyMsg(w dns.ResponseWriter, req *dns.Msg, err int) *dns.Msg {
	m := new(dns.Msg)

	m.SetReply(req)

	if err != 0 {
		m.SetRcode(req, err)
	}

	m.Authoritative = false
	m.RecursionAvailable = true

	return m
}

// Forwards a DNS request to the specified nameservers. On success, the
// original reply packet will be returned to the caller. On failure, a
// new packet will be returned with `RCODE` set to `SERVFAIL`.
// Even though the original `ResponseWriter` object is taken as an argument,
// this function does not send a reply to the client. Instead, the
// packet is returned for further processing by the caller.
func getServerReply(w dns.ResponseWriter, req *dns.Msg) *dns.Msg {
	if *verbose {
		log.Print("Forwarding ", req.Question[0].Name, "/", dns.Type(req.Question[0].Qtype).String())
	}

	// create new DNS client

	client := &dns.Client{Net: "udp", ReadTimeout: 4 * time.Second, WriteTimeout: 4 * time.Second, SingleInflight: true}

	if _, tcp := w.RemoteAddr().(*net.TCPAddr); tcp {
		client.Net = "tcp"
	}

	var r *dns.Msg
	var err error

	// loop through the specified nameservers and forward them the request

	// the request ID is used as a starting point in order to introduce at least
	// some element of randomness, instead of always hitting the first nameserver

	for i := 0; i < len(nameservers); i++ {
		r, _, err = client.Exchange(req, nameservers[(int(req.Id) + i) % len(nameservers)])

		if err == nil {
			r.Compress = true

			return r
		}
	}

	log.Print("Failed to forward request.", err)
	return getEmptyMsg(w, req, dns.RcodeServerFailure)
}

// Handles an incoming DNS request packet. This function decides whether
// the hostname listed in the DNS packet is worthy of manipulation, or
// not. The IP addresses listed in the reply to the user for a marked
// hostname are added to the routing table at this time before a
// reply is sent back to the user, otherwise the user agent of the client
// might connect faster than the routing changes can be made.
func handleRequest(w dns.ResponseWriter, req *dns.Msg) {
	var m *dns.Msg

	if len(req.Question) > 0 && (req.Question[0].Name == "netflix.com." || strings.HasSuffix(req.Question[0].Name, ".netflix.com.")) {
		if req.Question[0].Qtype == dns.TypeA {
			m = getServerReply(w, req)
			for _, ans := range m.Answer {
				if ans.Header().Rrtype == dns.TypeA {
					ip := ans.(*dns.A).A.String()
					routedv4[ip] = struct{}{}

					log.Print("Re-routing ", ip, " for ", ans.Header().Name, "/", dns.Type(ans.Header().Rrtype).String())

					runAndLog(router, "add", ip + "/32", *gateway4)
				} else if ans.Header().Rrtype == dns.TypeAAAA {
					// sanity check for now, shouldn't happen afaik
					log.Print("WARNING: AAAA response in ", ans.Header().Name, "/A")
				}
			}
		} else if req.Question[0].Qtype == dns.TypeAAAA {
			if routev6 {
				m = getServerReply(w, req)
				for _, ans := range m.Answer {
					if ans.Header().Rrtype == dns.TypeAAAA {
						ip := ans.(*dns.AAAA).AAAA.String()
						routedv6[ip] = struct{}{}

						log.Print("Re-routing ", ip, " for ", ans.Header().Name, "/", dns.Type(ans.Header().Rrtype).String())

						runAndLog(router, "add", ip + "/128", *gateway6)
					} else if ans.Header().Rrtype == dns.TypeA {
						log.Print("WARNING: A response in ", ans.Header().Name, "/AAAA")
					}
				}
			} else {
				if *verbose {
					log.Print("Hijacking ", req.Question[0].Name, "/", dns.Type(req.Question[0].Qtype).String())
				}
				m = getEmptyMsg(w, req, 0)
			}
		} else {
			m = getServerReply(w, req)
		}
	} else {
		m = getServerReply(w, req)
	}

	w.WriteMsg(m)
}

// Removes the routed from the Windows Routing Table that have been
// added durin the lifetime of the server. Failure to call this function
// during exit may result in the inaccessibility of the added IP addresses.
func removeRoutes() {
	// remove IPv4 routes

	if len(routedv4) > 0 {
		log.Print("Removing routes...")

		for ip, _ := range routedv4 {
			runAndLog(router, "delete", ip + "/32")
		}
	}

	// remove IPv6 routes

	if routev6 && len(routedv6) > 0 {
		log.Print("Removing IPv6 routes...")

		for ip, _ := range routedv6 {
			runAndLog(router, "delete", ip + "/128")
		}
	}
}

// Main entry point of the application. Its behaviour can be modified
// via command line arguments as shown by the `flag` calls inside.
func main() {
	// set-up flags

	gateway4 = flag.String("r", "", "IPv4 gateway address for routing destination")
	gateway6 = flag.String("r6", "", "IPv6 gateway address for routing destination")
	dnsr1    = flag.String("dp", "8.8.8.8", "primary DNS server")
	dnsr2    = flag.String("ds", "8.8.4.4", "secondary DNS server")
	verbose  = flag.Bool("v", false, "verbose logging")

	flag.Usage = func() {
		flag.PrintDefaults()
	}

	flag.Parse()

	// find the route utility in %PATH%

	router, _ = exec.LookPath("route")
	if len(router) < 1 {
		log.Fatal("Unable to find the `route` command in your %PATH%.")
	}

	// read gateway from arguments

	if len(*gateway4) < 1 {
		log.Fatal("A gateway IP must be specified via argument `r`.")
	}

	if ip := net.ParseIP(*gateway4); ip != nil && ip.To4() != nil {
		*gateway4 = ip.String()
	} else {
		log.Fatal("Specified gateway IP is not a valid IPv4 address.")
	}

	// IPv6 gateway is optional

	if len(*gateway6) > 1 {
		if ip := net.ParseIP(*gateway6); ip != nil && ip.To4() == nil {
			*gateway6 = ip.String()
			  routev6 = true
		} else {
			log.Fatal("Specified gateway IP is not a valid IPv6 address.")
		}
	} else {
		routev6 = false
	}

	// allocate set for routed IP addresses

	routedv4 = make(map[string]struct{})

	if (routev6) {
		routedv6 = make(map[string]struct{})
	}

	// read DNS servers to forward to

	if ip := net.ParseIP(*dnsr1); ip != nil {
		*dnsr1 = ip.String()
	} else {
		log.Fatal("Specified primary DNS server is not a valid IP address.")
	}

	if ip := net.ParseIP(*dnsr2); ip != nil {
		*dnsr2 = ip.String()
	} else {
		log.Fatal("Specified secondary DNS server is not a valid IP address.")
	}

	nameservers = []string{*dnsr1 + ":53", *dnsr2 + ":53"}

	// start DNS server

	log.Print("Starting DNS resolver...")
	log.Print("Forwarding to ", nameservers)

	dns.HandleFunc(".", handleRequest)

	// start local UDP listener

	go func() {
		srv := &dns.Server{Addr: ":53", Net: "udp"}
		err := srv.ListenAndServe()

		if err != nil {
			log.Fatal("Failed to start UDP server.", err.Error())
		}
	}()

	// start local TCP listener

	go func() {
		srv := &dns.Server{Addr: ":53", Net: "tcp"}
		err := srv.ListenAndServe()

		if err != nil {
			log.Fatal("Failed to start TCP server.", err.Error())
		}
	}()

	// start listening for OS signals

	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)

	for {
		select {
		case s := <-sig:
			removeRoutes()
			log.Fatalf("Received signal %d, exiting...", s)
		}
	}
}<|endoftext|>


<commit_before>package pipe

import (
	"io"
	"os"
	"sync"

	"github.com/CodisLabs/codis/pkg/utils/errors"
)

type Pipe struct {
	rd, wt struct {
		sync.Mutex
		cond *sync.Cond
		err  error
	}
	mu sync.Mutex

	store Buffer
}

func NewPipe() *Pipe {
	return NewPipeSize(defaultMemBufferSize)
}

func NewPipeSize(size int) *Pipe {
	return newPipe(newMemBufferSize(size))
}

func NewPipeFile(file *os.File, size int) *Pipe {
	return newPipe(newFileBufferSize(file, size))
}

func newPipe(store Buffer) *Pipe {
	p := &Pipe{store: store}
	p.rd.cond = sync.NewCond(&p.mu)
	p.wt.cond = sync.NewCond(&p.mu)
	return p
}

func (p *Pipe) Close() {
	p.CloseReader(nil)
	p.CloseWriter(nil)
}

func (p *Pipe) Reader() Reader {
	return &PipeReader{p}
}

func (p *Pipe) Read(b []byte) (int, error) {
	p.rd.Lock()
	defer p.rd.Unlock()
	for {
		n, err := p.readSome(b)
		if err != nil || n != 0 {
			return n, err
		}
		if len(b) == 0 {
			return 0, nil
		}
	}
}

func (p *Pipe) readSome(b []byte) (int, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.rd.err != nil {
		return 0, errors.Trace(io.ErrClosedPipe)
	}
	if len(b) == 0 {
		if p.store.Buffered() != 0 {
			return 0, nil
		}
		return 0, p.wt.err
	}
	n, err := p.store.ReadSome(b)
	if err != nil || n != 0 {
		p.wt.cond.Signal()
		return n, err
	}
	if p.wt.err != nil {
		return 0, p.wt.err
	} else {
		p.rd.cond.Wait()
		return 0, nil
	}
}

func (p *Pipe) Buffered() (int, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.rd.err != nil {
		return 0, p.rd.err
	}
	if n := p.store.Buffered(); n != 0 {
		return n, nil
	} else {
		return 0, p.wt.err
	}
}

func (p *Pipe) CloseReader(err error) error {
	if err == nil {
		err = errors.Trace(io.ErrClosedPipe)
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.rd.err == nil {
		p.rd.err = err
	}
	p.rd.cond.Broadcast()
	p.wt.cond.Broadcast()
	return p.store.CloseReader()
}

func (p *Pipe) Writer() Writer {
	return &PipeWriter{p}
}

func (p *Pipe) Write(b []byte) (int, error) {
	p.wt.Lock()
	defer p.wt.Unlock()
	var nn int
	for {
		n, err := p.writeSome(b)
		if err != nil || n == len(b) {
			return nn + n, err
		}
		nn, b = nn+n, b[n:]
	}
}

func (p *Pipe) writeSome(b []byte) (int, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.wt.err != nil {
		return 0, errors.Trace(io.ErrClosedPipe)
	}
	if p.rd.err != nil {
		return 0, p.rd.err
	}
	if len(b) == 0 {
		return 0, nil
	}
	n, err := p.store.WriteSome(b)
	if err != nil || n != 0 {
		p.rd.cond.Signal()
		return n, err
	} else {
		p.wt.cond.Wait()
		return 0, nil
	}
}

func (p *Pipe) Available() (int, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.wt.err != nil {
		return 0, p.wt.err
	}
	if p.rd.err != nil {
		return 0, p.rd.err
	}
	return p.store.Available(), nil
}

func (p *Pipe) CloseWriter(err error) error {
	if err == nil {
		err = errors.Trace(io.EOF)
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.wt.err == nil {
		p.wt.err = err
	}
	p.rd.cond.Broadcast()
	p.wt.cond.Broadcast()
	return p.store.CloseWriter()
}
<commit_msg>pipe: refactor<commit_after>package pipe

import (
	"io"
	"os"
	"sync"

	"github.com/CodisLabs/codis/pkg/utils/errors"
)

type Pipe struct {
	rd, wt struct {
		sync.Mutex
		cond *sync.Cond
		err  error
	}
	mu sync.Mutex

	store Buffer
}

func NewPipe() *Pipe {
	return NewPipeSize(defaultMemBufferSize)
}

func NewPipeSize(size int) *Pipe {
	return newPipe(newMemBufferSize(size))
}

func NewPipeFile(file *os.File, size int) *Pipe {
	return newPipe(newFileBufferSize(file, size))
}

func newPipe(store Buffer) *Pipe {
	p := &Pipe{store: store}
	p.rd.cond = sync.NewCond(&p.mu)
	p.wt.cond = sync.NewCond(&p.mu)
	return p
}

func (p *Pipe) Close() {
	p.CloseReader(nil)
	p.CloseWriter(nil)
}

func (p *Pipe) Reader() Reader {
	return &PipeReader{p}
}

func (p *Pipe) Read(b []byte) (int, error) {
	p.rd.Lock()
	defer p.rd.Unlock()
	p.mu.Lock()
	defer p.mu.Unlock()
	for {
		if p.rd.err != nil {
			return 0, errors.Trace(io.ErrClosedPipe)
		}
		if len(b) == 0 {
			if p.store.Buffered() != 0 {
				return 0, nil
			}
			return 0, p.wt.err
		}
		n, err := p.store.ReadSome(b)
		if err != nil || n != 0 {
			p.wt.cond.Signal()
			return n, err
		}
		if p.wt.err != nil {
			return 0, p.wt.err
		}
		p.rd.cond.Wait()
	}
}

func (p *Pipe) Buffered() (int, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.rd.err != nil {
		return 0, p.rd.err
	}
	if n := p.store.Buffered(); n != 0 {
		return n, nil
	} else {
		return 0, p.wt.err
	}
}

func (p *Pipe) CloseReader(err error) error {
	if err == nil {
		err = errors.Trace(io.ErrClosedPipe)
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.rd.err == nil {
		p.rd.err = err
	}
	p.rd.cond.Broadcast()
	p.wt.cond.Broadcast()
	return p.store.CloseReader()
}

func (p *Pipe) Writer() Writer {
	return &PipeWriter{p}
}

func (p *Pipe) Write(b []byte) (int, error) {
	p.wt.Lock()
	defer p.wt.Unlock()
	p.mu.Lock()
	defer p.mu.Unlock()
	var nn int
	for {
		if p.wt.err != nil {
			return nn, errors.Trace(io.ErrClosedPipe)
		}
		if p.rd.err != nil {
			return nn, p.rd.err
		}
		if len(b) == 0 {
			return nn, nil
		}
	again:
		n, err := p.store.WriteSome(b)
		if err != nil || n != 0 {
			p.rd.cond.Signal()
			nn, b = nn+n, b[n:]
			if err == nil && len(b) != 0 {
				goto again
			}
			return nn, err
		}
		p.wt.cond.Wait()
	}
}

func (p *Pipe) Available() (int, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.wt.err != nil {
		return 0, p.wt.err
	}
	if p.rd.err != nil {
		return 0, p.rd.err
	}
	return p.store.Available(), nil
}

func (p *Pipe) CloseWriter(err error) error {
	if err == nil {
		err = errors.Trace(io.EOF)
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.wt.err == nil {
		p.wt.err = err
	}
	p.rd.cond.Broadcast()
	p.wt.cond.Broadcast()
	return p.store.CloseWriter()
}
<|endoftext|>


<commit_before>// +build !openbsd
// +build !freebsd
// +build !netbsd

package logger

import (
	"os"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/shirou/gopsutil/process"
)

func GetLoggingLevelByInspectingParent() (logrus.Level, error) {
	ppid := os.Getppid()
	process, err := process.NewProcess(int32(ppid))
	if err != nil {
		return logrus.WarnLevel, err
	}

	cmdline, err := process.Cmdline()
	if err != nil {
		return logrus.WarnLevel, err
	}

	if strings.Contains(cmdline, "-vv") {
		return logrus.DebugLevel, nil
	} else if strings.Contains(cmdline, "-v") {
		return logrus.InfoLevel, nil
	}
	return logrus.WarnLevel, nil
}
<commit_msg>Add quiet option check<commit_after>// +build !openbsd
// +build !freebsd
// +build !netbsd

package logger

import (
	"os"
	"strings"

	"github.com/Sirupsen/logrus"
	"github.com/shirou/gopsutil/process"
)

func GetLoggingLevelByInspectingParent() (logrus.Level, error) {
	ppid := os.Getppid()
	process, err := process.NewProcess(int32(ppid))
	if err != nil {
		return logrus.WarnLevel, err
	}

	cmdline, err := process.Cmdline()
	if err != nil {
		return logrus.WarnLevel, err
	}

	if strings.Contains(cmdline, "-vv") {
		return logrus.DebugLevel, nil
	} else if strings.Contains(cmdline, "-v") {
		return logrus.InfoLevel, nil
	} else if strings.Contains(cmdline, "-q") {
		return logrus.ErrorLevel, nil
	}
	return logrus.WarnLevel, nil
}
<|endoftext|>


<commit_before>/*
Copyright The Helm 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 strvals

import (
	"bytes"
	"fmt"
	"io"
	"strconv"
	"strings"

	"github.com/pkg/errors"
	"sigs.k8s.io/yaml"
)

// ErrNotList indicates that a non-list was treated as a list.
var ErrNotList = errors.New("not a list")

// ToYAML takes a string of arguments and converts to a YAML document.
func ToYAML(s string) (string, error) {
	m, err := Parse(s)
	if err != nil {
		return "", err
	}
	d, err := yaml.Marshal(m)
	return strings.TrimSuffix(string(d), "\n"), err
}

// Parse parses a set line.
//
// A set line is of the form name1=value1,name2=value2
func Parse(s string) (map[string]interface{}, error) {
	vals := map[string]interface{}{}
	scanner := bytes.NewBufferString(s)
	t := newParser(scanner, vals, false)
	err := t.parse()
	return vals, err
}

// ParseString parses a set line and forces a string value.
//
// A set line is of the form name1=value1,name2=value2
func ParseString(s string) (map[string]interface{}, error) {
	vals := map[string]interface{}{}
	scanner := bytes.NewBufferString(s)
	t := newParser(scanner, vals, true)
	err := t.parse()
	return vals, err
}

// ParseInto parses a strvals line and merges the result into dest.
//
// If the strval string has a key that exists in dest, it overwrites the
// dest version.
func ParseInto(s string, dest map[string]interface{}) error {
	scanner := bytes.NewBufferString(s)
	t := newParser(scanner, dest, false)
	return t.parse()
}

// ParseFile parses a set line, but its final value is loaded from the file at the path specified by the original value.
//
// A set line is of the form name1=path1,name2=path2
//
// When the files at path1 and path2 contained "val1" and "val2" respectively, the set line is consumed as
// name1=val1,name2=val2
func ParseFile(s string, reader RunesValueReader) (map[string]interface{}, error) {
	vals := map[string]interface{}{}
	scanner := bytes.NewBufferString(s)
	t := newFileParser(scanner, vals, reader)
	err := t.parse()
	return vals, err
}

// ParseIntoString parses a strvals line and merges the result into dest.
//
// This method always returns a string as the value.
func ParseIntoString(s string, dest map[string]interface{}) error {
	scanner := bytes.NewBufferString(s)
	t := newParser(scanner, dest, true)
	return t.parse()
}

// ParseIntoFile parses a filevals line and merges the result into dest.
//
// This method always returns a string as the value.
func ParseIntoFile(s string, dest map[string]interface{}, reader RunesValueReader) error {
	scanner := bytes.NewBufferString(s)
	t := newFileParser(scanner, dest, reader)
	return t.parse()
}

// RunesValueReader is a function that takes the given value (a slice of runes)
// and returns the parsed value
type RunesValueReader func([]rune) (interface{}, error)

// parser is a simple parser that takes a strvals line and parses it into a
// map representation.
//
// where sc is the source of the original data being parsed
// where data is the final parsed data from the parses with correct types
type parser struct {
	sc     *bytes.Buffer
	data   map[string]interface{}
	reader RunesValueReader
}

func newParser(sc *bytes.Buffer, data map[string]interface{}, stringBool bool) *parser {
	stringConverter := func(rs []rune) (interface{}, error) {
		return typedVal(rs, stringBool), nil
	}
	return &parser{sc: sc, data: data, reader: stringConverter}
}

func newFileParser(sc *bytes.Buffer, data map[string]interface{}, reader RunesValueReader) *parser {
	return &parser{sc: sc, data: data, reader: reader}
}

func (t *parser) parse() error {
	for {
		err := t.key(t.data)
		if err == nil {
			continue
		}
		if err == io.EOF {
			return nil
		}
		return err
	}
}

func runeSet(r []rune) map[rune]bool {
	s := make(map[rune]bool, len(r))
	for _, rr := range r {
		s[rr] = true
	}
	return s
}

func (t *parser) key(data map[string]interface{}) (reterr error) {
	defer func() {
		if r := recover(); r != nil {
			reterr = fmt.Errorf("unable to parse key: %s", r)
		}
	}()
	stop := runeSet([]rune{'=', '[', ',', '.'})
	for {
		switch k, last, err := runesUntil(t.sc, stop); {
		case err != nil:
			if len(k) == 0 {
				return err
			}
			return errors.Errorf("key %q has no value", string(k))
			//set(data, string(k), "")
			//return err
		case last == '[':
			// We are in a list index context, so we need to set an index.
			i, err := t.keyIndex()
			if err != nil {
				return errors.Wrap(err, "error parsing index")
			}
			kk := string(k)
			// Find or create target list
			list := []interface{}{}
			if _, ok := data[kk]; ok {
				list = data[kk].([]interface{})
			}

			// Now we need to get the value after the ].
			list, err = t.listItem(list, i)
			set(data, kk, list)
			return err
		case last == '=':
			//End of key. Consume =, Get value.
			// FIXME: Get value list first
			vl, e := t.valList()
			switch e {
			case nil:
				set(data, string(k), vl)
				return nil
			case io.EOF:
				set(data, string(k), "")
				return e
			case ErrNotList:
				rs, e := t.val()
				if e != nil && e != io.EOF {
					return e
				}
				v, e := t.reader(rs)
				set(data, string(k), v)
				return e
			default:
				return e
			}

		case last == ',':
			// No value given. Set the value to empty string. Return error.
			set(data, string(k), "")
			return errors.Errorf("key %q has no value (cannot end with ,)", string(k))
		case last == '.':
			// First, create or find the target map.
			inner := map[string]interface{}{}
			if _, ok := data[string(k)]; ok {
				inner = data[string(k)].(map[string]interface{})
			}

			// Recurse
			e := t.key(inner)
			if len(inner) == 0 {
				return errors.Errorf("key map %q has no value", string(k))
			}
			set(data, string(k), inner)
			return e
		}
	}
}

func set(data map[string]interface{}, key string, val interface{}) {
	// If key is empty, don't set it.
	if len(key) == 0 {
		return
	}
	data[key] = val
}

func setIndex(list []interface{}, index int, val interface{}) (l2 []interface{}, err error) {
	// There are possible index values that are out of range on a target system
	// causing a panic. This will catch the panic and return an error instead.
	// The value of the index that causes a panic varies from system to system.
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("error processing index %d: %s", index, r)
		}
	}()

	if index < 0 {
		return list, fmt.Errorf("negative %d index not allowed", index)
	}
	if len(list) <= index {
		newlist := make([]interface{}, index+1)
		copy(newlist, list)
		list = newlist
	}
	list[index] = val
	return list, nil
}

func (t *parser) keyIndex() (int, error) {
	// First, get the key.
	stop := runeSet([]rune{']'})
	v, _, err := runesUntil(t.sc, stop)
	if err != nil {
		return 0, err
	}
	// v should be the index
	return strconv.Atoi(string(v))

}
func (t *parser) listItem(list []interface{}, i int) ([]interface{}, error) {
	if i < 0 {
		return list, fmt.Errorf("negative %d index not allowed", i)
	}
	stop := runeSet([]rune{'[', '.', '='})
	switch k, last, err := runesUntil(t.sc, stop); {
	case len(k) > 0:
		return list, errors.Errorf("unexpected data at end of array index: %q", k)
	case err != nil:
		return list, err
	case last == '=':
		vl, e := t.valList()
		switch e {
		case nil:
			return setIndex(list, i, vl)
		case io.EOF:
			return setIndex(list, i, "")
		case ErrNotList:
			rs, e := t.val()
			if e != nil && e != io.EOF {
				return list, e
			}
			v, e := t.reader(rs)
			if e != nil {
				return list, e
			}
			return setIndex(list, i, v)
		default:
			return list, e
		}
	case last == '[':
		// now we have a nested list. Read the index and handle.
		nextI, err := t.keyIndex()
		if err != nil {
			return list, errors.Wrap(err, "error parsing index")
		}
		var crtList []interface{}
		if len(list) > i {
			// If nested list already exists, take the value of list to next cycle.
			existed := list[i]
			if existed != nil {
				crtList = list[i].([]interface{})
			}
		}
		// Now we need to get the value after the ].
<<<<<<< HEAD
		list2, err := t.listItem(crtList, i)
		if err != nil {
			return list, err
		}
		return setIndex(list, i, list2)
=======
		list2, err := t.listItem(crtList, nextI)
		return setIndex(list, i, list2), err
>>>>>>> fix another extreme case
	case last == '.':
		// We have a nested object. Send to t.key
		inner := map[string]interface{}{}
		if len(list) > i {
			var ok bool
			inner, ok = list[i].(map[string]interface{})
			if !ok {
				// We have indices out of order. Initialize empty value.
				list[i] = map[string]interface{}{}
				inner = list[i].(map[string]interface{})
			}
		}

		// Recurse
		e := t.key(inner)
		if e != nil {
			return list, e
		}
		return setIndex(list, i, inner)
	default:
		return nil, errors.Errorf("parse error: unexpected token %v", last)
	}
}

func (t *parser) val() ([]rune, error) {
	stop := runeSet([]rune{','})
	v, _, err := runesUntil(t.sc, stop)
	return v, err
}

func (t *parser) valList() ([]interface{}, error) {
	r, _, e := t.sc.ReadRune()
	if e != nil {
		return []interface{}{}, e
	}

	if r != '{' {
		t.sc.UnreadRune()
		return []interface{}{}, ErrNotList
	}

	list := []interface{}{}
	stop := runeSet([]rune{',', '}'})
	for {
		switch rs, last, err := runesUntil(t.sc, stop); {
		case err != nil:
			if err == io.EOF {
				err = errors.New("list must terminate with '}'")
			}
			return list, err
		case last == '}':
			// If this is followed by ',', consume it.
			if r, _, e := t.sc.ReadRune(); e == nil && r != ',' {
				t.sc.UnreadRune()
			}
			v, e := t.reader(rs)
			list = append(list, v)
			return list, e
		case last == ',':
			v, e := t.reader(rs)
			if e != nil {
				return list, e
			}
			list = append(list, v)
		}
	}
}

func runesUntil(in io.RuneReader, stop map[rune]bool) ([]rune, rune, error) {
	v := []rune{}
	for {
		switch r, _, e := in.ReadRune(); {
		case e != nil:
			return v, r, e
		case inMap(r, stop):
			return v, r, nil
		case r == '\\':
			next, _, e := in.ReadRune()
			if e != nil {
				return v, next, e
			}
			v = append(v, next)
		default:
			v = append(v, r)
		}
	}
}

func inMap(k rune, m map[rune]bool) bool {
	_, ok := m[k]
	return ok
}

func typedVal(v []rune, st bool) interface{} {
	val := string(v)

	if st {
		return val
	}

	if strings.EqualFold(val, "true") {
		return true
	}

	if strings.EqualFold(val, "false") {
		return false
	}

	if strings.EqualFold(val, "null") {
		return nil
	}

	if strings.EqualFold(val, "0") {
		return int64(0)
	}

	// If this value does not start with zero, try parsing it to an int
	if len(val) != 0 && val[0] != '0' {
		if iv, err := strconv.ParseInt(val, 10, 64); err == nil {
			return iv
		}
	}

	return val
}
<commit_msg>fix conflict<commit_after>/*
Copyright The Helm 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 strvals

import (
	"bytes"
	"fmt"
	"io"
	"strconv"
	"strings"

	"github.com/pkg/errors"
	"sigs.k8s.io/yaml"
)

// ErrNotList indicates that a non-list was treated as a list.
var ErrNotList = errors.New("not a list")

// ToYAML takes a string of arguments and converts to a YAML document.
func ToYAML(s string) (string, error) {
	m, err := Parse(s)
	if err != nil {
		return "", err
	}
	d, err := yaml.Marshal(m)
	return strings.TrimSuffix(string(d), "\n"), err
}

// Parse parses a set line.
//
// A set line is of the form name1=value1,name2=value2
func Parse(s string) (map[string]interface{}, error) {
	vals := map[string]interface{}{}
	scanner := bytes.NewBufferString(s)
	t := newParser(scanner, vals, false)
	err := t.parse()
	return vals, err
}

// ParseString parses a set line and forces a string value.
//
// A set line is of the form name1=value1,name2=value2
func ParseString(s string) (map[string]interface{}, error) {
	vals := map[string]interface{}{}
	scanner := bytes.NewBufferString(s)
	t := newParser(scanner, vals, true)
	err := t.parse()
	return vals, err
}

// ParseInto parses a strvals line and merges the result into dest.
//
// If the strval string has a key that exists in dest, it overwrites the
// dest version.
func ParseInto(s string, dest map[string]interface{}) error {
	scanner := bytes.NewBufferString(s)
	t := newParser(scanner, dest, false)
	return t.parse()
}

// ParseFile parses a set line, but its final value is loaded from the file at the path specified by the original value.
//
// A set line is of the form name1=path1,name2=path2
//
// When the files at path1 and path2 contained "val1" and "val2" respectively, the set line is consumed as
// name1=val1,name2=val2
func ParseFile(s string, reader RunesValueReader) (map[string]interface{}, error) {
	vals := map[string]interface{}{}
	scanner := bytes.NewBufferString(s)
	t := newFileParser(scanner, vals, reader)
	err := t.parse()
	return vals, err
}

// ParseIntoString parses a strvals line and merges the result into dest.
//
// This method always returns a string as the value.
func ParseIntoString(s string, dest map[string]interface{}) error {
	scanner := bytes.NewBufferString(s)
	t := newParser(scanner, dest, true)
	return t.parse()
}

// ParseIntoFile parses a filevals line and merges the result into dest.
//
// This method always returns a string as the value.
func ParseIntoFile(s string, dest map[string]interface{}, reader RunesValueReader) error {
	scanner := bytes.NewBufferString(s)
	t := newFileParser(scanner, dest, reader)
	return t.parse()
}

// RunesValueReader is a function that takes the given value (a slice of runes)
// and returns the parsed value
type RunesValueReader func([]rune) (interface{}, error)

// parser is a simple parser that takes a strvals line and parses it into a
// map representation.
//
// where sc is the source of the original data being parsed
// where data is the final parsed data from the parses with correct types
type parser struct {
	sc     *bytes.Buffer
	data   map[string]interface{}
	reader RunesValueReader
}

func newParser(sc *bytes.Buffer, data map[string]interface{}, stringBool bool) *parser {
	stringConverter := func(rs []rune) (interface{}, error) {
		return typedVal(rs, stringBool), nil
	}
	return &parser{sc: sc, data: data, reader: stringConverter}
}

func newFileParser(sc *bytes.Buffer, data map[string]interface{}, reader RunesValueReader) *parser {
	return &parser{sc: sc, data: data, reader: reader}
}

func (t *parser) parse() error {
	for {
		err := t.key(t.data)
		if err == nil {
			continue
		}
		if err == io.EOF {
			return nil
		}
		return err
	}
}

func runeSet(r []rune) map[rune]bool {
	s := make(map[rune]bool, len(r))
	for _, rr := range r {
		s[rr] = true
	}
	return s
}

func (t *parser) key(data map[string]interface{}) (reterr error) {
	defer func() {
		if r := recover(); r != nil {
			reterr = fmt.Errorf("unable to parse key: %s", r)
		}
	}()
	stop := runeSet([]rune{'=', '[', ',', '.'})
	for {
		switch k, last, err := runesUntil(t.sc, stop); {
		case err != nil:
			if len(k) == 0 {
				return err
			}
			return errors.Errorf("key %q has no value", string(k))
			//set(data, string(k), "")
			//return err
		case last == '[':
			// We are in a list index context, so we need to set an index.
			i, err := t.keyIndex()
			if err != nil {
				return errors.Wrap(err, "error parsing index")
			}
			kk := string(k)
			// Find or create target list
			list := []interface{}{}
			if _, ok := data[kk]; ok {
				list = data[kk].([]interface{})
			}

			// Now we need to get the value after the ].
			list, err = t.listItem(list, i)
			set(data, kk, list)
			return err
		case last == '=':
			//End of key. Consume =, Get value.
			// FIXME: Get value list first
			vl, e := t.valList()
			switch e {
			case nil:
				set(data, string(k), vl)
				return nil
			case io.EOF:
				set(data, string(k), "")
				return e
			case ErrNotList:
				rs, e := t.val()
				if e != nil && e != io.EOF {
					return e
				}
				v, e := t.reader(rs)
				set(data, string(k), v)
				return e
			default:
				return e
			}

		case last == ',':
			// No value given. Set the value to empty string. Return error.
			set(data, string(k), "")
			return errors.Errorf("key %q has no value (cannot end with ,)", string(k))
		case last == '.':
			// First, create or find the target map.
			inner := map[string]interface{}{}
			if _, ok := data[string(k)]; ok {
				inner = data[string(k)].(map[string]interface{})
			}

			// Recurse
			e := t.key(inner)
			if len(inner) == 0 {
				return errors.Errorf("key map %q has no value", string(k))
			}
			set(data, string(k), inner)
			return e
		}
	}
}

func set(data map[string]interface{}, key string, val interface{}) {
	// If key is empty, don't set it.
	if len(key) == 0 {
		return
	}
	data[key] = val
}

func setIndex(list []interface{}, index int, val interface{}) (l2 []interface{}, err error) {
	// There are possible index values that are out of range on a target system
	// causing a panic. This will catch the panic and return an error instead.
	// The value of the index that causes a panic varies from system to system.
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("error processing index %d: %s", index, r)
		}
	}()

	if index < 0 {
		return list, fmt.Errorf("negative %d index not allowed", index)
	}
	if len(list) <= index {
		newlist := make([]interface{}, index+1)
		copy(newlist, list)
		list = newlist
	}
	list[index] = val
	return list, nil
}

func (t *parser) keyIndex() (int, error) {
	// First, get the key.
	stop := runeSet([]rune{']'})
	v, _, err := runesUntil(t.sc, stop)
	if err != nil {
		return 0, err
	}
	// v should be the index
	return strconv.Atoi(string(v))

}
func (t *parser) listItem(list []interface{}, i int) ([]interface{}, error) {
	if i < 0 {
		return list, fmt.Errorf("negative %d index not allowed", i)
	}
	stop := runeSet([]rune{'[', '.', '='})
	switch k, last, err := runesUntil(t.sc, stop); {
	case len(k) > 0:
		return list, errors.Errorf("unexpected data at end of array index: %q", k)
	case err != nil:
		return list, err
	case last == '=':
		vl, e := t.valList()
		switch e {
		case nil:
			return setIndex(list, i, vl)
		case io.EOF:
			return setIndex(list, i, "")
		case ErrNotList:
			rs, e := t.val()
			if e != nil && e != io.EOF {
				return list, e
			}
			v, e := t.reader(rs)
			if e != nil {
				return list, e
			}
			return setIndex(list, i, v)
		default:
			return list, e
		}
	case last == '[':
		// now we have a nested list. Read the index and handle.
		nextI, err := t.keyIndex()
		if err != nil {
			return list, errors.Wrap(err, "error parsing index")
		}
		var crtList []interface{}
		if len(list) > i {
			// If nested list already exists, take the value of list to next cycle.
			existed := list[i]
			if existed != nil {
				crtList = list[i].([]interface{})
			}
		}
		// Now we need to get the value after the ].
		list2, err := t.listItem(crtList, nextI)
		if err != nil {
			return list, err
		}
		return setIndex(list, i, list2)
	case last == '.':
		// We have a nested object. Send to t.key
		inner := map[string]interface{}{}
		if len(list) > i {
			var ok bool
			inner, ok = list[i].(map[string]interface{})
			if !ok {
				// We have indices out of order. Initialize empty value.
				list[i] = map[string]interface{}{}
				inner = list[i].(map[string]interface{})
			}
		}

		// Recurse
		e := t.key(inner)
		if e != nil {
			return list, e
		}
		return setIndex(list, i, inner)
	default:
		return nil, errors.Errorf("parse error: unexpected token %v", last)
	}
}

func (t *parser) val() ([]rune, error) {
	stop := runeSet([]rune{','})
	v, _, err := runesUntil(t.sc, stop)
	return v, err
}

func (t *parser) valList() ([]interface{}, error) {
	r, _, e := t.sc.ReadRune()
	if e != nil {
		return []interface{}{}, e
	}

	if r != '{' {
		t.sc.UnreadRune()
		return []interface{}{}, ErrNotList
	}

	list := []interface{}{}
	stop := runeSet([]rune{',', '}'})
	for {
		switch rs, last, err := runesUntil(t.sc, stop); {
		case err != nil:
			if err == io.EOF {
				err = errors.New("list must terminate with '}'")
			}
			return list, err
		case last == '}':
			// If this is followed by ',', consume it.
			if r, _, e := t.sc.ReadRune(); e == nil && r != ',' {
				t.sc.UnreadRune()
			}
			v, e := t.reader(rs)
			list = append(list, v)
			return list, e
		case last == ',':
			v, e := t.reader(rs)
			if e != nil {
				return list, e
			}
			list = append(list, v)
		}
	}
}

func runesUntil(in io.RuneReader, stop map[rune]bool) ([]rune, rune, error) {
	v := []rune{}
	for {
		switch r, _, e := in.ReadRune(); {
		case e != nil:
			return v, r, e
		case inMap(r, stop):
			return v, r, nil
		case r == '\\':
			next, _, e := in.ReadRune()
			if e != nil {
				return v, next, e
			}
			v = append(v, next)
		default:
			v = append(v, r)
		}
	}
}

func inMap(k rune, m map[rune]bool) bool {
	_, ok := m[k]
	return ok
}

func typedVal(v []rune, st bool) interface{} {
	val := string(v)

	if st {
		return val
	}

	if strings.EqualFold(val, "true") {
		return true
	}

	if strings.EqualFold(val, "false") {
		return false
	}

	if strings.EqualFold(val, "null") {
		return nil
	}

	if strings.EqualFold(val, "0") {
		return int64(0)
	}

	// If this value does not start with zero, try parsing it to an int
	if len(val) != 0 && val[0] != '0' {
		if iv, err := strconv.ParseInt(val, 10, 64); err == nil {
			return iv
		}
	}

	return val
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The CDI 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 system

import (
	"bufio"
	"bytes"
	"context"
	"os/exec"
	"syscall"
	"time"
	"unsafe"

	"github.com/pkg/errors"
	"golang.org/x/sys/unix"
	"k8s.io/klog"
)

// ProcessLimiter defines the methods limiting resources of a Process
type ProcessLimiter interface {
	SetAddressSpaceLimit(pid int, value uint64) error
	SetCPUTimeLimit(pid int, value uint64) error
}

// ProcessLimitValues specifies the resource limits available to a process
type ProcessLimitValues struct {
	AddressSpaceLimit uint64
	CPUTimeLimit      uint64
}

type processLimiter struct{}

var execCommand = exec.Command
var execCommandContext = exec.CommandContext

var limiter = NewProcessLimiter()

// NewProcessLimiter returns a new ProcessLimiter
func NewProcessLimiter() ProcessLimiter {
	return &processLimiter{}
}

func (p *processLimiter) SetAddressSpaceLimit(pid int, value uint64) error {
	return prlimit(pid, unix.RLIMIT_AS, &syscall.Rlimit{Cur: value, Max: value})
}

func (p *processLimiter) SetCPUTimeLimit(pid int, value uint64) error {
	return prlimit(pid, unix.RLIMIT_CPU, &syscall.Rlimit{Cur: value, Max: value})
}

// SetAddressSpaceLimit sets a limit on total address space of a process
func SetAddressSpaceLimit(pid int, value uint64) error {
	return limiter.SetAddressSpaceLimit(pid, value)
}

// SetCPUTimeLimit sets a limit on the total cpu time a process may have
func SetCPUTimeLimit(pid int, value uint64) error {
	return limiter.SetCPUTimeLimit(pid, value)
}

// scanLinesWithCR is an alternate split function that works with carriage returns as well
// as new lines.
func scanLinesWithCR(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}
	if i := bytes.IndexByte(data, '\r'); i >= 0 {
		// We have a full carriage return-terminated line.
		return i + 1, data[0:i], nil
	}
	if i := bytes.IndexByte(data, '\n'); i >= 0 {
		// We have a full newline-terminated line.
		return i + 1, data[0:i], nil
	}
	// If we're at EOF, we have a final, non-terminated line. Return it.
	if atEOF {
		return len(data), data, nil
	}
	// Request more data.
	return 0, nil, nil
}

func processScanner(scanner *bufio.Scanner, buf *bytes.Buffer, done chan bool, callback func(string)) {
	for scanner.Scan() {
		line := scanner.Text()
		buf.WriteString(line)
		if callback != nil {
			callback(line)
		}
	}
	done <- true
}

// ExecWithLimits executes a command with process limits
func ExecWithLimits(limits *ProcessLimitValues, callback func(string), command string, args ...string) ([]byte, error) {
	// Args can potentially contain sensitive information, make sure NOT to write args to the logs.
	var buf bytes.Buffer
	var cmd *exec.Cmd

	stdoutDone := make(chan bool)
	stderrDone := make(chan bool)

	if limits != nil && limits.CPUTimeLimit > 0 {
		klog.V(3).Infof("Setting CPU limit to %d\n", limits.CPUTimeLimit)
		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(limits.CPUTimeLimit)*time.Second)
		defer cancel()
		cmd = execCommandContext(ctx, command, args...)
	} else {
		cmd = execCommand(command, args...)
	}
	stdoutIn, err := cmd.StdoutPipe()
	if err != nil {
		return nil, errors.Wrapf(err, "Couldn't get stdout for %s", command)
	}
	stderrIn, err := cmd.StderrPipe()
	if err != nil {
		return nil, errors.Wrapf(err, "Couldn't get stderr for %s", command)
	}

	scanner := bufio.NewScanner(stdoutIn)
	scanner.Split(scanLinesWithCR)
	errScanner := bufio.NewScanner(stderrIn)
	errScanner.Split(scanLinesWithCR)

	err = cmd.Start()
	if err != nil {
		return nil, errors.Wrapf(err, "Couldn't start %s", command)
	}
	defer cmd.Process.Kill()

	go processScanner(scanner, &buf, stdoutDone, callback)
	go processScanner(errScanner, &buf, stderrDone, callback)

	if limits != nil && limits.AddressSpaceLimit > 0 {
		klog.V(3).Infof("Setting Address space limit to %d\n", limits.AddressSpaceLimit)
		err = SetAddressSpaceLimit(cmd.Process.Pid, limits.AddressSpaceLimit)
		if err != nil {
			return nil, errors.Wrap(err, "Couldn't set address space limit")
		}
	}

	err = cmd.Wait()
	<-stdoutDone
	<-stderrDone

	output := buf.Bytes()
	if err != nil {
		klog.Errorf("%s failed output is:\n", command)
		klog.Errorf("%s\n", string(output))
		return output, errors.Wrapf(err, "%s execution failed", command)
	}
	return output, nil
}

func prlimit(pid int, limit int, value *syscall.Rlimit) error {
	_, _, e1 := syscall.RawSyscall6(syscall.SYS_PRLIMIT64, uintptr(pid), uintptr(limit), uintptr(unsafe.Pointer(value)), 0, 0, 0)
	if e1 != 0 {
		return errors.Wrapf(e1, "error setting prlimit on %d with value %d on pid %d", limit, value, pid)
	}
	return nil
}
<commit_msg>Fix race in cmd.Exec where sometimes stdout/err was closed before being read. (#936)<commit_after>/*
Copyright 2018 The CDI 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 system

import (
	"bufio"
	"bytes"
	"context"
	"os/exec"
	"syscall"
	"time"
	"unsafe"

	"github.com/pkg/errors"
	"golang.org/x/sys/unix"
	"k8s.io/klog"
)

// ProcessLimiter defines the methods limiting resources of a Process
type ProcessLimiter interface {
	SetAddressSpaceLimit(pid int, value uint64) error
	SetCPUTimeLimit(pid int, value uint64) error
}

// ProcessLimitValues specifies the resource limits available to a process
type ProcessLimitValues struct {
	AddressSpaceLimit uint64
	CPUTimeLimit      uint64
}

type processLimiter struct{}

var execCommand = exec.Command
var execCommandContext = exec.CommandContext

var limiter = NewProcessLimiter()

// NewProcessLimiter returns a new ProcessLimiter
func NewProcessLimiter() ProcessLimiter {
	return &processLimiter{}
}

func (p *processLimiter) SetAddressSpaceLimit(pid int, value uint64) error {
	return prlimit(pid, unix.RLIMIT_AS, &syscall.Rlimit{Cur: value, Max: value})
}

func (p *processLimiter) SetCPUTimeLimit(pid int, value uint64) error {
	return prlimit(pid, unix.RLIMIT_CPU, &syscall.Rlimit{Cur: value, Max: value})
}

// SetAddressSpaceLimit sets a limit on total address space of a process
func SetAddressSpaceLimit(pid int, value uint64) error {
	return limiter.SetAddressSpaceLimit(pid, value)
}

// SetCPUTimeLimit sets a limit on the total cpu time a process may have
func SetCPUTimeLimit(pid int, value uint64) error {
	return limiter.SetCPUTimeLimit(pid, value)
}

// scanLinesWithCR is an alternate split function that works with carriage returns as well
// as new lines.
func scanLinesWithCR(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}
	if i := bytes.IndexByte(data, '\r'); i >= 0 {
		// We have a full carriage return-terminated line.
		return i + 1, data[0:i], nil
	}
	if i := bytes.IndexByte(data, '\n'); i >= 0 {
		// We have a full newline-terminated line.
		return i + 1, data[0:i], nil
	}
	// If we're at EOF, we have a final, non-terminated line. Return it.
	if atEOF {
		return len(data), data, nil
	}
	// Request more data.
	return 0, nil, nil
}

func processScanner(scanner *bufio.Scanner, buf *bytes.Buffer, done chan bool, callback func(string)) {
	for scanner.Scan() {
		line := scanner.Text()
		buf.WriteString(line)
		if callback != nil {
			callback(line)
		}
	}
	done <- true
}

// ExecWithLimits executes a command with process limits
func ExecWithLimits(limits *ProcessLimitValues, callback func(string), command string, args ...string) ([]byte, error) {
	// Args can potentially contain sensitive information, make sure NOT to write args to the logs.
	var buf bytes.Buffer
	var cmd *exec.Cmd

	stdoutDone := make(chan bool)
	stderrDone := make(chan bool)

	if limits != nil && limits.CPUTimeLimit > 0 {
		klog.V(3).Infof("Setting CPU limit to %d\n", limits.CPUTimeLimit)
		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(limits.CPUTimeLimit)*time.Second)
		defer cancel()
		cmd = execCommandContext(ctx, command, args...)
	} else {
		cmd = execCommand(command, args...)
	}
	stdoutIn, err := cmd.StdoutPipe()
	if err != nil {
		return nil, errors.Wrapf(err, "Couldn't get stdout for %s", command)
	}
	stderrIn, err := cmd.StderrPipe()
	if err != nil {
		return nil, errors.Wrapf(err, "Couldn't get stderr for %s", command)
	}

	scanner := bufio.NewScanner(stdoutIn)
	scanner.Split(scanLinesWithCR)
	errScanner := bufio.NewScanner(stderrIn)
	errScanner.Split(scanLinesWithCR)

	err = cmd.Start()
	if err != nil {
		return nil, errors.Wrapf(err, "Couldn't start %s", command)
	}
	defer cmd.Process.Kill()

	go processScanner(scanner, &buf, stdoutDone, callback)
	go processScanner(errScanner, &buf, stderrDone, callback)

	if limits != nil && limits.AddressSpaceLimit > 0 {
		klog.V(3).Infof("Setting Address space limit to %d\n", limits.AddressSpaceLimit)
		err = SetAddressSpaceLimit(cmd.Process.Pid, limits.AddressSpaceLimit)
		if err != nil {
			return nil, errors.Wrap(err, "Couldn't set address space limit")
		}
	}
	<-stdoutDone
	<-stderrDone
	// The wait has to be after the reading channels are finished otherwise there is a race where the wait completes and closes stdout/err before anything
	// is read from it.
	err = cmd.Wait()

	output := buf.Bytes()
	if err != nil {
		klog.Errorf("%s failed output is:\n", command)
		klog.Errorf("%s\n", string(output))
		return output, errors.Wrapf(err, "%s execution failed", command)
	}
	return output, nil
}

func prlimit(pid int, limit int, value *syscall.Rlimit) error {
	_, _, e1 := syscall.RawSyscall6(syscall.SYS_PRLIMIT64, uintptr(pid), uintptr(limit), uintptr(unsafe.Pointer(value)), 0, 0, 0)
	if e1 != 0 {
		return errors.Wrapf(e1, "error setting prlimit on %d with value %d on pid %d", limit, value, pid)
	}
	return nil
}
<|endoftext|>