<commit_before>// +build linux

package cpu

import (
	"errors"
	"strconv"
	"strings"

	common "github.com/shirou/gopsutil/common"
)

func CPUTimes(percpu bool) ([]CPUTimesStat, error) {
	filename := "/proc/stat"
	var lines = []string{}
	if percpu {
		var startIdx uint = 1
		for {
			linen, _ := common.ReadLinesOffsetN(filename, startIdx, 1)
			line := linen[0]
			if !strings.HasPrefix(line, "cpu") {
				break
			}
			lines = append(lines, line)
			startIdx += 1
		}
	} else {
		lines, _ = common.ReadLinesOffsetN(filename, 0, 1)
	}

	ret := make([]CPUTimesStat, 0, len(lines))

	for _, line := range lines {
		ct, err := parseStatLine(line)
		if err != nil {
			continue
		}
		ret = append(ret, *ct)

	}
	return ret, nil
}

func CPUInfo() ([]CPUInfoStat, error) {
	filename := "/proc/cpuinfo"
	lines, _ := common.ReadLines(filename)

	var ret []CPUInfoStat

	var c CPUInfoStat
	for _, line := range lines {
		fields := strings.Split(line, ":")
		if len(fields) < 2 {
			if c.VendorID != "" {
				ret = append(ret, c)
			}
			continue
		}
		key := strings.TrimSpace(fields[0])
		value := strings.TrimSpace(fields[1])

		switch key {
		case "processor":
			c = CPUInfoStat{}
			t, err := strconv.ParseInt(value, 10, 64)
			if err != nil {
				return ret, err
			}
			c.CPU = int32(t)
		case "vendor_id":
			c.VendorID = value
		case "cpu family":
			c.Family = value
		case "model":
			c.Model = value
		case "model name":
			c.ModelName = value
		case "stepping":
			t, err := strconv.ParseInt(value, 10, 64)
			if err != nil {
				return ret, err
			}
			c.Stepping = int32(t)
		case "cpu MHz":
			t, err := strconv.ParseFloat(value, 64)
			if err != nil {
				return ret, err
			}
			c.Mhz = t
		case "cache size":
			t, err := strconv.ParseInt(strings.Replace(value, " KB", "", 1), 10, 64)
			if err != nil {
				return ret, err
			}
			c.CacheSize = int32(t)
		case "physical id":
			c.PhysicalID = value
		case "core id":
			c.CoreID = value
		case "cpu cores":
			t, err := strconv.ParseInt(value, 10, 64)
			if err != nil {
				return ret, err
			}
			c.Cores = int32(t)
		case "flags":
			c.Flags = strings.Split(value, ",")
		}
	}
	return ret, nil
}

func parseStatLine(line string) (*CPUTimesStat, error) {
	fields := strings.Fields(line)

	if strings.HasPrefix(fields[0], "cpu") == false {
		//		return CPUTimesStat{}, e
		return nil, errors.New("not contain cpu")
	}

	cpu := fields[0]
	if cpu == "cpu" {
		cpu = "cpu-total"
	}
	user, err := strconv.ParseFloat(fields[1], 64)
	if err != nil {
		return nil, err
	}
	nice, err := strconv.ParseFloat(fields[2], 64)
	if err != nil {
		return nil, err
	}
	system, err := strconv.ParseFloat(fields[3], 64)
	if err != nil {
		return nil, err
	}
	idle, err := strconv.ParseFloat(fields[4], 64)
	if err != nil {
		return nil, err
	}
	iowait, err := strconv.ParseFloat(fields[5], 64)
	if err != nil {
		return nil, err
	}
	irq, err := strconv.ParseFloat(fields[6], 64)
	if err != nil {
		return nil, err
	}
	softirq, err := strconv.ParseFloat(fields[7], 64)
	if err != nil {
		return nil, err
	}
	stolen, err := strconv.ParseFloat(fields[8], 64)
	if err != nil {
		return nil, err
	}

	cpu_tick := float64(100) // TODO: how to get _SC_CLK_TCK ?
	ct := &CPUTimesStat{
		CPU:     cpu,
		User:    float64(user) / cpu_tick,
		Nice:    float64(nice) / cpu_tick,
		System:  float64(system) / cpu_tick,
		Idle:    float64(idle) / cpu_tick,
		Iowait:  float64(iowait) / cpu_tick,
		Irq:     float64(irq) / cpu_tick,
		Softirq: float64(softirq) / cpu_tick,
		Stolen:  float64(stolen) / cpu_tick,
	}
	if len(fields) > 9 { // Linux >= 2.6.11
		steal, err := strconv.ParseFloat(fields[9], 64)
		if err != nil {
			return nil, err
		}
		ct.Steal = float64(steal)
	}
	if len(fields) > 10 { // Linux >= 2.6.24
		guest, err := strconv.ParseFloat(fields[10], 64)
		if err != nil {
			return nil, err
		}
		ct.Guest = float64(guest)
	}
	if len(fields) > 11 { // Linux >= 3.2.0
		guestNice, err := strconv.ParseFloat(fields[11], 64)
		if err != nil {
			return nil, err
		}
		ct.GuestNice = float64(guestNice)
	}

	return ct, nil
}
<commit_msg>Fix interpretation error about /proc/stat<commit_after>// +build linux

package cpu

import (
	"errors"
	"strconv"
	"strings"

	common "github.com/shirou/gopsutil/common"
)

func CPUTimes(percpu bool) ([]CPUTimesStat, error) {
	filename := "/proc/stat"
	var lines = []string{}
	if percpu {
		var startIdx uint = 1
		for {
			linen, _ := common.ReadLinesOffsetN(filename, startIdx, 1)
			line := linen[0]
			if !strings.HasPrefix(line, "cpu") {
				break
			}
			lines = append(lines, line)
			startIdx += 1
		}
	} else {
		lines, _ = common.ReadLinesOffsetN(filename, 0, 1)
	}

	ret := make([]CPUTimesStat, 0, len(lines))

	for _, line := range lines {
		ct, err := parseStatLine(line)
		if err != nil {
			continue
		}
		ret = append(ret, *ct)

	}
	return ret, nil
}

func CPUInfo() ([]CPUInfoStat, error) {
	filename := "/proc/cpuinfo"
	lines, _ := common.ReadLines(filename)

	var ret []CPUInfoStat

	var c CPUInfoStat
	for _, line := range lines {
		fields := strings.Split(line, ":")
		if len(fields) < 2 {
			if c.VendorID != "" {
				ret = append(ret, c)
			}
			continue
		}
		key := strings.TrimSpace(fields[0])
		value := strings.TrimSpace(fields[1])

		switch key {
		case "processor":
			c = CPUInfoStat{}
			t, err := strconv.ParseInt(value, 10, 64)
			if err != nil {
				return ret, err
			}
			c.CPU = int32(t)
		case "vendor_id":
			c.VendorID = value
		case "cpu family":
			c.Family = value
		case "model":
			c.Model = value
		case "model name":
			c.ModelName = value
		case "stepping":
			t, err := strconv.ParseInt(value, 10, 64)
			if err != nil {
				return ret, err
			}
			c.Stepping = int32(t)
		case "cpu MHz":
			t, err := strconv.ParseFloat(value, 64)
			if err != nil {
				return ret, err
			}
			c.Mhz = t
		case "cache size":
			t, err := strconv.ParseInt(strings.Replace(value, " KB", "", 1), 10, 64)
			if err != nil {
				return ret, err
			}
			c.CacheSize = int32(t)
		case "physical id":
			c.PhysicalID = value
		case "core id":
			c.CoreID = value
		case "cpu cores":
			t, err := strconv.ParseInt(value, 10, 64)
			if err != nil {
				return ret, err
			}
			c.Cores = int32(t)
		case "flags":
			c.Flags = strings.Split(value, ",")
		}
	}
	return ret, nil
}

func parseStatLine(line string) (*CPUTimesStat, error) {
	fields := strings.Fields(line)

	if strings.HasPrefix(fields[0], "cpu") == false {
		//		return CPUTimesStat{}, e
		return nil, errors.New("not contain cpu")
	}

	cpu := fields[0]
	if cpu == "cpu" {
		cpu = "cpu-total"
	}
	user, err := strconv.ParseFloat(fields[1], 64)
	if err != nil {
		return nil, err
	}
	nice, err := strconv.ParseFloat(fields[2], 64)
	if err != nil {
		return nil, err
	}
	system, err := strconv.ParseFloat(fields[3], 64)
	if err != nil {
		return nil, err
	}
	idle, err := strconv.ParseFloat(fields[4], 64)
	if err != nil {
		return nil, err
	}
	iowait, err := strconv.ParseFloat(fields[5], 64)
	if err != nil {
		return nil, err
	}
	irq, err := strconv.ParseFloat(fields[6], 64)
	if err != nil {
		return nil, err
	}
	softirq, err := strconv.ParseFloat(fields[7], 64)
	if err != nil {
		return nil, err
	}

	cpu_tick := float64(100) // TODO: how to get _SC_CLK_TCK ?
	ct := &CPUTimesStat{
		CPU:     cpu,
		User:    float64(user) / cpu_tick,
		Nice:    float64(nice) / cpu_tick,
		System:  float64(system) / cpu_tick,
		Idle:    float64(idle) / cpu_tick,
		Iowait:  float64(iowait) / cpu_tick,
		Irq:     float64(irq) / cpu_tick,
		Softirq: float64(softirq) / cpu_tick,
	}
	if len(fields) > 8 { // Linux >= 2.6.11
		steal, err := strconv.ParseFloat(fields[8], 64)
		if err != nil {
			return nil, err
		}
		ct.Steal = float64(steal)
	}
	if len(fields) > 9 { // Linux >= 2.6.24
		guest, err := strconv.ParseFloat(fields[9], 64)
		if err != nil {
			return nil, err
		}
		ct.Guest = float64(guest)
	}
	if len(fields) > 10 { // Linux >= 3.2.0
		guestNice, err := strconv.ParseFloat(fields[10], 64)
		if err != nil {
			return nil, err
		}
		ct.GuestNice = float64(guestNice)
	}

	return ct, nil
}
<|endoftext|>


<commit_before>package cpu

import (
	"bufio"
	"os"
	"regexp"
)

var cpuMap = map[string]string{
	"vendor_id":  "vendor_id",
	"model name": "model_name",
	"cpu cores":  "cpu_cores",
	"siblings":   "cpu_logical_processors",
	"cpu MHz\t":  "mhz",
	"cache size": "cache_size",
	"cpu family": "family",
	"model\t":    "model",
	"stepping":   "stepping",
}

func getCpuInfo() (cpuInfo map[string]string, err error) {
	file, err := os.Open("/proc/cpuinfo")

	if err != nil {
		return
	}

	var lines []string
	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}

	if scanner.Err() != nil {
		err = scanner.Err()
		return
	}

	cpuInfo = make(map[string]string)

	for _, line := range lines[1:] {
		pair := regexp.MustCompile("\t: ").Split(line, 2)

		key, ok := cpuMap[pair[0]]
		if ok {
			cpuInfo[key] = pair[1]
		}
	}

	return
}
<commit_msg>[cpu] Take into account environments with multiple physical procs (#32)<commit_after>package cpu

import (
	"bufio"
	"os"
	"regexp"
	"strconv"
)

var cpuMap = map[string]string{
	"vendor_id":  "vendor_id",
	"model name": "model_name",
	"cpu cores":  "cpu_cores",
	"siblings":   "cpu_logical_processors",
	"cpu MHz\t":  "mhz",
	"cache size": "cache_size",
	"cpu family": "family",
	"model\t":    "model",
	"stepping":   "stepping",
}

// Values that need to be multiplied by the number of physical processors
var perPhysicalProcValues = []string{
	"cpu_cores",
	"cpu_logical_processors",
}

func getCpuInfo() (cpuInfo map[string]string, err error) {
	lines, err := readProcFile()
	if err != nil {
		return
	}

	cpuInfo = make(map[string]string)
	// Implementation of a set that holds the physical IDs
	physicalProcIDs := make(map[string]struct{})

	for _, line := range lines {
		pair := regexp.MustCompile("\t: ").Split(line, 2)

		if pair[0] == "physical id" {
			physicalProcIDs[pair[1]] = struct{}{}
		}

		key, ok := cpuMap[pair[0]]
		if ok {
			cpuInfo[key] = pair[1]
		}
	}

	// Multiply the values that are "per physical processor" by the number of physical procs
	for _, field := range perPhysicalProcValues {
		if value, ok := cpuInfo[field]; ok {
			intValue, err := strconv.Atoi(value)
			if err != nil {
				continue
			}

			cpuInfo[field] = strconv.Itoa(intValue * len(physicalProcIDs))
		}
	}

	return
}

func readProcFile() (lines []string, err error) {
	file, err := os.Open("/proc/cpuinfo")

	if err != nil {
		return
	}

	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}

	if scanner.Err() != nil {
		err = scanner.Err()
		return
	}

	return
}
<|endoftext|>


<commit_before>package fsql

import (
	"bytes"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"testing"
	"time"
)

var files = map[string]*os.FileInfo{}

func TestRun_All(t *testing.T) {
	type Case struct {
		query    string
		expected string
	}

	cases := []Case{
		{
			query: "SELECT all FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("drwxr-xr-x\t%s\t-------\tfoo\n",
				strings.Join(GetAttrs("foo", "size", "time"), "\t")),
		},
		{
			query: "SELECT all FROM ./testdata WHERE name LIKE waldo",
			expected: fmt.Sprintf("-rw-r--r--\t%s\twaldo\n",
				strings.Join(GetAttrs("foo/quuz/waldo", "size", "time", "hash"), "\t")),
		},
		{
			query:    "SELECT all FROM ./testdata WHERE FORMAT(time, 'Jan 02 2006 15:04') > 'Jan 01 2999 00:00'",
			expected: "",
		},
		{
			query: "SELECT all FROM ./testdata WHERE mode IS DIR",
			expected: fmt.Sprintf(
				strings.Repeat("%s\n", 8),
				fmt.Sprintf(
					"drwxr-xr-x\t%s\t-------\t%-8s",
					strings.Join(GetAttrs(".", "size", "time"), "\t"),
					"testdata",
				),
				fmt.Sprintf(
					"drwxr-xr-x\t%s\t-------\t%-8s",
					strings.Join(GetAttrs("bar", "size", "time"), "\t"),
					"bar",
				),
				fmt.Sprintf(
					"drwxr-xr-x\t%s\t-------\t%-8s",
					strings.Join(GetAttrs("bar/garply", "size", "time"), "\t"),
					"garply",
				),
				fmt.Sprintf(
					"drwxr-xr-x\t%s\t-------\t%-8s",
					strings.Join(GetAttrs("bar/garply/xyzzy", "size", "time"), "\t"),
					"xyzzy",
				),
				fmt.Sprintf(
					"drwxr-xr-x\t%s\t-------\t%-8s",
					strings.Join(GetAttrs("bar/garply/xyzzy/thud", "size", "time"), "\t"),
					"thud",
				),
				fmt.Sprintf(
					"drwxr-xr-x\t%s\t-------\t%-8s",
					strings.Join(GetAttrs("foo", "size", "time"), "\t"),
					"foo",
				),
				fmt.Sprintf(
					"drwxr-xr-x\t%s\t-------\t%-8s",
					strings.Join(GetAttrs("foo/quuz", "size", "time"), "\t"),
					"quuz",
				),
				fmt.Sprintf(
					"drwxr-xr-x\t%s\t-------\t%-8s",
					strings.Join(GetAttrs("foo/quuz/fred", "size", "time"), "\t"),
					"fred",
				),
			),
		},
	}

	for _, c := range cases {
		actual := DoRun(c.query)
		if !reflect.DeepEqual(c.expected, actual) {
			t.Fatalf("\nExpected:\n%v\nGot:\n%v", c.expected, actual)
		}
	}
}

func TestRun_Multiple(t *testing.T) {
	type Case struct {
		query    string
		expected string
	}

	cases := []Case{
		{
			query:    "SELECT name, size FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("foo\t%s\n", GetAttrs("foo", "size")[0]),
		},
		{
			query:    "SELECT size, name FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("%s\tfoo\n", GetAttrs("foo", "size")[0]),
		},
		{
			query: "SELECT FULLPATH(name), size, time FROM ./testdata/foo",
			expected: fmt.Sprintf(
				strings.Repeat("%s\n", 7),
				fmt.Sprintf(
					"%-31s\t%s",
					"testdata/foo",
					strings.Join(GetAttrs("foo", "size", "time"), "\t"),
				),
				fmt.Sprintf(
					"%-31s\t%s",
					"testdata/foo/quux",
					strings.Join(GetAttrs("foo/quux", "size", "time"), "\t"),
				),
				fmt.Sprintf(
					"%-31s\t%s",
					"testdata/foo/quuz",
					strings.Join(GetAttrs("foo/quuz", "size", "time"), "\t"),
				),
				fmt.Sprintf(
					"%-31s\t%s",
					"testdata/foo/quuz/fred",
					strings.Join(GetAttrs("foo/quuz/fred", "size", "time"), "\t"),
				),
				fmt.Sprintf(
					"%-31s\t%s",
					"testdata/foo/quuz/fred/.gitkeep",
					strings.Join(GetAttrs("foo/quuz/fred/.gitkeep", "size", "time"), "\t"),
				),
				fmt.Sprintf(
					"%-31s\t%s",
					"testdata/foo/quuz/waldo",
					strings.Join(GetAttrs("foo/quuz/waldo", "size", "time"), "\t"),
				),
				fmt.Sprintf(
					"%-31s\t%s",
					"testdata/foo/qux",
					strings.Join(GetAttrs("foo/qux", "size", "time"), "\t"),
				),
			),
		},
	}

	for _, c := range cases {
		actual := DoRun(c.query)
		if !reflect.DeepEqual(c.expected, actual) {
			t.Fatalf("\nExpected:\n%v\nGot:\n%v", c.expected, actual)
		}
	}
}

func TestRun_Name(t *testing.T) {
	type Case struct {
		query    string
		expected string
	}

	cases := []Case{
		{
			query:    "SELECT name FROM ./testdata WHERE name REGEXP ^g.*",
			expected: "garply\ngrault\n",
		},
		{
			query:    "SELECT FULLPATH(name) FROM ./testdata WHERE name REGEXP ^b.*",
			expected: "testdata/bar\ntestdata/baz\n",
		},
		{
			query: "SELECT UPPER(FULLPATH(name)) FROM ./testdata WHERE mode IS DIR",
			expected: fmt.Sprintf(
				strings.Repeat("%-30s\n", 8),
				"TESTDATA",
				"TESTDATA/BAR",
				"TESTDATA/BAR/GARPLY",
				"TESTDATA/BAR/GARPLY/XYZZY",
				"TESTDATA/BAR/GARPLY/XYZZY/THUD",
				"TESTDATA/FOO",
				"TESTDATA/FOO/QUUZ",
				"TESTDATA/FOO/QUUZ/FRED",
			),
		},
	}

	for _, c := range cases {
		actual := DoRun(c.query)
		if !reflect.DeepEqual(c.expected, actual) {
			t.Fatalf("\nExpected:\n%v\nGot:\n%v", c.expected, actual)
		}
	}
}

func TestRun_Size(t *testing.T) {
	type Case struct {
		query    string
		expected string
	}

	cases := []Case{
		{
			query:    "SELECT size FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("%s\n", GetAttrs("foo", "size")[0]),
		},
		{
			query:    "SELECT FORMAT(size, KB) FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("%s\n", GetAttrs("foo", "size:kb")[0]),
		},
		{
			query:    "SELECT FORMAT(size, MB) FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("%s\n", GetAttrs("foo", "size:mb")[0]),
		},
		{
			query:    "SELECT FORMAT(size, GB) FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("%s\n", GetAttrs("foo", "size:gb")[0]),
		},
		{
			query: "SELECT size FROM ./testdata WHERE name LIKE qu",
			expected: fmt.Sprintf(
				strings.Repeat("%s\n", 3),
				GetAttrs("foo/quux", "size")[0],
				GetAttrs("foo/quuz", "size")[0],
				GetAttrs("foo/qux", "size")[0],
			),
		},
	}

	for _, c := range cases {
		actual := DoRun(c.query)
		if !reflect.DeepEqual(c.expected, actual) {
			t.Fatalf("\nExpected:\n%v\nGot:\n%v", c.expected, actual)
		}
	}
}

func TestRun_Time(t *testing.T) {
	type Case struct {
		query    string
		expected string
	}

	cases := []Case{
		{
			query:    "SELECT time FROM ./testdata WHERE name = baz",
			expected: fmt.Sprintf("%s\n", GetAttrs("baz", "time")[0]),
		},
		{
			query:    "SELECT FORMAT(time, ISO) FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("%s\n", GetAttrs("foo", "time:iso")[0]),
		},
		{
			query:    "SELECT FORMAT(time, 2006) FROM ./testdata WHERE NOT name LIKE .%",
			expected: strings.Repeat(fmt.Sprintf("%s\n", GetAttrs(".", "time:year")[0]), 14),
		},
		{
			query: "SELECT time FROM ./testdata/foo/quuz",
			expected: fmt.Sprintf(
				strings.Repeat("%s\n", 4),
				GetAttrs("foo/quuz", "time")[0],
				GetAttrs("foo/quuz/fred", "time")[0],
				GetAttrs("foo/quuz/fred/.gitkeep", "time")[0],
				GetAttrs("foo/quuz/waldo", "time")[0],
			),
		},
	}

	for _, c := range cases {
		actual := DoRun(c.query)
		if !reflect.DeepEqual(c.expected, actual) {
			t.Fatalf("\nExpected:\n%v\nGot:\n%v", c.expected, actual)
		}
	}
}

func TestRun_Mode(t *testing.T) {
	type Case struct {
		query    string
		expected string
	}

	cases := []Case{
		{
			query:    "SELECT mode FROM ./testdata WHERE name = foo",
			expected: "drwxr-xr-x\n",
		},
		{
			query:    "SELECT mode FROM ./testdata WHERE name = baz",
			expected: "-rwxr-xr-x\n",
		},
		{
			query:    "SELECT mode FROM ./testdata WHERE mode IS DIR",
			expected: strings.Repeat("drwxr-xr-x\n", 8),
		},
	}

	for _, c := range cases {
		actual := DoRun(c.query)
		if !reflect.DeepEqual(c.expected, actual) {
			t.Fatalf("\nExpected:\n%v\nGot:\n%v", c.expected, actual)
		}
	}
}

func TestRun_Hash(t *testing.T) {
	type Case struct {
		query    string
		expected string
	}

	// TODO
	cases := []Case{}

	for _, c := range cases {
		actual := DoRun(c.query)
		if !reflect.DeepEqual(c.expected, actual) {
			t.Fatalf("\nExpected:\n%v\nGot:\n%v", c.expected, actual)
		}
	}
}

func GetAttrs(path string, attrs ...string) []string {
	// If the files map is empty, walk ./testdata and populate it.
	if len(files) == 0 {
		if err := filepath.Walk(
			"./testdata",
			func(path string, info os.FileInfo, err error) error {
				files[filepath.Clean(path)] = &info
				return nil
			},
		); err != nil {
			return []string{}
		}
	}

	path = filepath.Clean(fmt.Sprintf("testdata/%s", path))
	file, ok := files[path]
	if !ok {
		return []string{}
	}

	result := make([]string, len(attrs))
	for i, attr := range attrs {
		// Hard-coding modifiers works for the time being, but we might need a more
		// elegant solution when we introduce new modifiers in the future.
		switch attr {
		case "hash":
			b, err := ioutil.ReadFile(path)
			if err != nil {
				return []string{}
			}
			h := sha1.New()
			if _, err := h.Write(b); err != nil {
				return []string{}
			}
			result[i] = hex.EncodeToString(h.Sum(nil))[:7]
		case "size":
			result[i] = fmt.Sprintf("%d", (*file).Size())
		case "size:kb", "size:mb", "size:gb":
			size := (*file).Size()
			switch attr[len(attr)-2:] {
			case "kb":
				result[i] = fmt.Sprintf("%fkb", float64(size)/(1<<10))
			case "mb":
				result[i] = fmt.Sprintf("%fmb", float64(size)/(1<<20))
			case "gb":
				result[i] = fmt.Sprintf("%fgb", float64(size)/(1<<30))
			}
		case "time":
			result[i] = (*file).ModTime().Format(time.Stamp)
		case "time:iso":
			result[i] = (*file).ModTime().Format(time.RFC3339)
		case "time:year":
			result[i] = (*file).ModTime().Format("2006")
		}
	}
	return result
}

// DoRun executes fsql.Run and returns the output.
func DoRun(query string) string {
	stdout := os.Stdout
	ch := make(chan string)

	r, w, err := os.Pipe()
	if err != nil {
		return ""
	}
	os.Stdout = w

	if err := Run(query); err != nil {
		return ""
	}

	go func() {
		var buf bytes.Buffer
		io.Copy(&buf, r)
		ch <- buf.String()
	}()

	w.Close()
	os.Stdout = stdout
	return <-ch
}
<commit_msg>Don't hardcode test directory mode strings<commit_after>package fsql

import (
	"bytes"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"testing"
	"time"
)

var files = map[string]*os.FileInfo{}

func TestRun_All(t *testing.T) {
	type Case struct {
		query    string
		expected string
	}

	cases := []Case{
		{
			query: "SELECT all FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("%s\t-------\tfoo\n",
				strings.Join(GetAttrs("foo", "mode", "size", "time"), "\t")),
		},
		{
			query: "SELECT all FROM ./testdata WHERE name LIKE waldo",
			expected: fmt.Sprintf("%s\twaldo\n",
				strings.Join(GetAttrs("foo/quuz/waldo", "mode", "size", "time", "hash"), "\t")),
		},
		{
			query:    "SELECT all FROM ./testdata WHERE FORMAT(time, 'Jan 02 2006 15:04') > 'Jan 01 2999 00:00'",
			expected: "",
		},
		{
			query: "SELECT all FROM ./testdata WHERE mode IS DIR",
			expected: fmt.Sprintf(
				strings.Repeat("%s\n", 8),
				fmt.Sprintf(
					"%s\t-------\t%-8s",
					strings.Join(GetAttrs(".", "mode", "size", "time"), "\t"),
					"testdata",
				),
				fmt.Sprintf(
					"%s\t-------\t%-8s",
					strings.Join(GetAttrs("bar", "mode", "size", "time"), "\t"),
					"bar",
				),
				fmt.Sprintf(
					"%s\t-------\t%-8s",
					strings.Join(GetAttrs("bar/garply", "mode", "size", "time"), "\t"),
					"garply",
				),
				fmt.Sprintf(
					"%s\t-------\t%-8s",
					strings.Join(GetAttrs("bar/garply/xyzzy", "mode", "size", "time"), "\t"),
					"xyzzy",
				),
				fmt.Sprintf(
					"%s\t-------\t%-8s",
					strings.Join(GetAttrs("bar/garply/xyzzy/thud", "mode", "size", "time"), "\t"),
					"thud",
				),
				fmt.Sprintf(
					"%s\t-------\t%-8s",
					strings.Join(GetAttrs("foo", "mode", "size", "time"), "\t"),
					"foo",
				),
				fmt.Sprintf(
					"%s\t-------\t%-8s",
					strings.Join(GetAttrs("foo/quuz", "mode", "size", "time"), "\t"),
					"quuz",
				),
				fmt.Sprintf(
					"%s\t-------\t%-8s",
					strings.Join(GetAttrs("foo/quuz/fred", "mode", "size", "time"), "\t"),
					"fred",
				),
			),
		},
	}

	for _, c := range cases {
		actual := DoRun(c.query)
		if !reflect.DeepEqual(c.expected, actual) {
			t.Fatalf("%s\nExpected:\n%v\nGot:\n%v", c.query, c.expected, actual)
		}
	}
}

func TestRun_Multiple(t *testing.T) {
	type Case struct {
		query    string
		expected string
	}

	cases := []Case{
		{
			query:    "SELECT name, size FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("foo\t%s\n", GetAttrs("foo", "size")[0]),
		},
		{
			query:    "SELECT size, name FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("%s\tfoo\n", GetAttrs("foo", "size")[0]),
		},
		{
			query: "SELECT FULLPATH(name), size, time FROM ./testdata/foo",
			expected: fmt.Sprintf(
				strings.Repeat("%s\n", 7),
				fmt.Sprintf(
					"%-31s\t%s",
					"testdata/foo",
					strings.Join(GetAttrs("foo", "size", "time"), "\t"),
				),
				fmt.Sprintf(
					"%-31s\t%s",
					"testdata/foo/quux",
					strings.Join(GetAttrs("foo/quux", "size", "time"), "\t"),
				),
				fmt.Sprintf(
					"%-31s\t%s",
					"testdata/foo/quuz",
					strings.Join(GetAttrs("foo/quuz", "size", "time"), "\t"),
				),
				fmt.Sprintf(
					"%-31s\t%s",
					"testdata/foo/quuz/fred",
					strings.Join(GetAttrs("foo/quuz/fred", "size", "time"), "\t"),
				),
				fmt.Sprintf(
					"%-31s\t%s",
					"testdata/foo/quuz/fred/.gitkeep",
					strings.Join(GetAttrs("foo/quuz/fred/.gitkeep", "size", "time"), "\t"),
				),
				fmt.Sprintf(
					"%-31s\t%s",
					"testdata/foo/quuz/waldo",
					strings.Join(GetAttrs("foo/quuz/waldo", "size", "time"), "\t"),
				),
				fmt.Sprintf(
					"%-31s\t%s",
					"testdata/foo/qux",
					strings.Join(GetAttrs("foo/qux", "size", "time"), "\t"),
				),
			),
		},
	}

	for _, c := range cases {
		actual := DoRun(c.query)
		if !reflect.DeepEqual(c.expected, actual) {
			t.Fatalf("%s\nExpected:\n%v\nGot:\n%v", c.query, c.expected, actual)
		}
	}
}

func TestRun_Name(t *testing.T) {
	type Case struct {
		query    string
		expected string
	}

	cases := []Case{
		{
			query:    "SELECT name FROM ./testdata WHERE name REGEXP ^g.*",
			expected: "garply\ngrault\n",
		},
		{
			query:    "SELECT FULLPATH(name) FROM ./testdata WHERE name REGEXP ^b.*",
			expected: "testdata/bar\ntestdata/baz\n",
		},
		{
			query: "SELECT UPPER(FULLPATH(name)) FROM ./testdata WHERE mode IS DIR",
			expected: fmt.Sprintf(
				strings.Repeat("%-30s\n", 8),
				"TESTDATA",
				"TESTDATA/BAR",
				"TESTDATA/BAR/GARPLY",
				"TESTDATA/BAR/GARPLY/XYZZY",
				"TESTDATA/BAR/GARPLY/XYZZY/THUD",
				"TESTDATA/FOO",
				"TESTDATA/FOO/QUUZ",
				"TESTDATA/FOO/QUUZ/FRED",
			),
		},
	}

	for _, c := range cases {
		actual := DoRun(c.query)
		if !reflect.DeepEqual(c.expected, actual) {
			t.Fatalf("%s\nExpected:\n%v\nGot:\n%v", c.query, c.expected, actual)
		}
	}
}

func TestRun_Size(t *testing.T) {
	type Case struct {
		query    string
		expected string
	}

	cases := []Case{
		{
			query:    "SELECT size FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("%s\n", GetAttrs("foo", "size")[0]),
		},
		{
			query:    "SELECT FORMAT(size, KB) FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("%s\n", GetAttrs("foo", "size:kb")[0]),
		},
		{
			query:    "SELECT FORMAT(size, MB) FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("%s\n", GetAttrs("foo", "size:mb")[0]),
		},
		{
			query:    "SELECT FORMAT(size, GB) FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("%s\n", GetAttrs("foo", "size:gb")[0]),
		},
		{
			query: "SELECT size FROM ./testdata WHERE name LIKE qu",
			expected: fmt.Sprintf(
				strings.Repeat("%s\n", 3),
				GetAttrs("foo/quux", "size")[0],
				GetAttrs("foo/quuz", "size")[0],
				GetAttrs("foo/qux", "size")[0],
			),
		},
	}

	for _, c := range cases {
		actual := DoRun(c.query)
		if !reflect.DeepEqual(c.expected, actual) {
			t.Fatalf("%s\nExpected:\n%v\nGot:\n%v", c.query, c.expected, actual)
		}
	}
}

func TestRun_Time(t *testing.T) {
	type Case struct {
		query    string
		expected string
	}

	cases := []Case{
		{
			query:    "SELECT time FROM ./testdata WHERE name = baz",
			expected: fmt.Sprintf("%s\n", GetAttrs("baz", "time")[0]),
		},
		{
			query:    "SELECT FORMAT(time, ISO) FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("%s\n", GetAttrs("foo", "time:iso")[0]),
		},
		{
			query:    "SELECT FORMAT(time, 2006) FROM ./testdata WHERE NOT name LIKE .%",
			expected: strings.Repeat(fmt.Sprintf("%s\n", GetAttrs(".", "time:year")[0]), 14),
		},
		{
			query: "SELECT time FROM ./testdata/foo/quuz",
			expected: fmt.Sprintf(
				strings.Repeat("%s\n", 4),
				GetAttrs("foo/quuz", "time")[0],
				GetAttrs("foo/quuz/fred", "time")[0],
				GetAttrs("foo/quuz/fred/.gitkeep", "time")[0],
				GetAttrs("foo/quuz/waldo", "time")[0],
			),
		},
	}

	for _, c := range cases {
		actual := DoRun(c.query)
		if !reflect.DeepEqual(c.expected, actual) {
			t.Fatalf("%s\nExpected:\n%v\nGot:\n%v", c.query, c.expected, actual)
		}
	}
}

func TestRun_Mode(t *testing.T) {
	type Case struct {
		query    string
		expected string
	}

	cases := []Case{
		{
			query:    "SELECT mode FROM ./testdata WHERE name = foo",
			expected: fmt.Sprintf("%s\n", GetAttrs("foo", "mode")[0]),
		},
		{
			query:    "SELECT mode FROM ./testdata WHERE name = baz",
			expected: fmt.Sprintf("%s\n", GetAttrs("baz", "mode")[0]),
		},
		{
			query: "SELECT mode FROM ./testdata WHERE mode IS DIR",
			expected: fmt.Sprintf(strings.Repeat("%s\n", 8),
				GetAttrs("", "mode")[0],
				GetAttrs("foo", "mode")[0],
				GetAttrs("foo/quuz", "mode")[0],
				GetAttrs("foo/quuz/fred", "mode")[0],
				GetAttrs("bar", "mode")[0],
				GetAttrs("bar/garply", "mode")[0],
				GetAttrs("bar/garply/xyzzy", "mode")[0],
				GetAttrs("bar/garply/xyzzy/thud", "mode")[0],
			),
		},
	}

	for _, c := range cases {
		actual := DoRun(c.query)
		if !reflect.DeepEqual(c.expected, actual) {
			t.Fatalf("%s\nExpected:\n\"%v\"\nGot:\n\"%v\"", c.query, c.expected, actual)
		}
	}
}

func TestRun_Hash(t *testing.T) {
	type Case struct {
		query    string
		expected string
	}

	// TODO
	cases := []Case{}

	for _, c := range cases {
		actual := DoRun(c.query)
		if !reflect.DeepEqual(c.expected, actual) {
			t.Fatalf("%s\nExpected:\n%v\nGot:\n%v", c.query, c.expected, actual)
		}
	}
}

func GetAttrs(path string, attrs ...string) []string {
	// If the files map is empty, walk ./testdata and populate it.
	if len(files) == 0 {
		if err := filepath.Walk(
			"./testdata",
			func(path string, info os.FileInfo, err error) error {
				files[filepath.Clean(path)] = &info
				return nil
			},
		); err != nil {
			return []string{}
		}
	}

	path = filepath.Clean(fmt.Sprintf("testdata/%s", path))
	file, ok := files[path]
	if !ok {
		return []string{}
	}

	result := make([]string, len(attrs))
	for i, attr := range attrs {
		// Hard-coding modifiers works for the time being, but we might need a more
		// elegant solution when we introduce new modifiers in the future.
		switch attr {
		case "hash":
			b, err := ioutil.ReadFile(path)
			if err != nil {
				return []string{}
			}
			h := sha1.New()
			if _, err := h.Write(b); err != nil {
				return []string{}
			}
			result[i] = hex.EncodeToString(h.Sum(nil))[:7]
		case "mode":
			result[i] = (*file).Mode().String()
		case "size":
			result[i] = fmt.Sprintf("%d", (*file).Size())
		case "size:kb", "size:mb", "size:gb":
			size := (*file).Size()
			switch attr[len(attr)-2:] {
			case "kb":
				result[i] = fmt.Sprintf("%fkb", float64(size)/(1<<10))
			case "mb":
				result[i] = fmt.Sprintf("%fmb", float64(size)/(1<<20))
			case "gb":
				result[i] = fmt.Sprintf("%fgb", float64(size)/(1<<30))
			}
		case "time":
			result[i] = (*file).ModTime().Format(time.Stamp)
		case "time:iso":
			result[i] = (*file).ModTime().Format(time.RFC3339)
		case "time:year":
			result[i] = (*file).ModTime().Format("2006")
		}
	}
	return result
}

// DoRun executes fsql.Run and returns the output.
func DoRun(query string) string {
	stdout := os.Stdout
	ch := make(chan string)

	r, w, err := os.Pipe()
	if err != nil {
		return ""
	}
	os.Stdout = w

	if err := Run(query); err != nil {
		return ""
	}

	go func() {
		var buf bytes.Buffer
		io.Copy(&buf, r)
		ch <- buf.String()
	}()

	w.Close()
	os.Stdout = stdout
	return <-ch
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
	"strings"

	log "github.com/sirupsen/logrus"
)

// CSPReport is the structure of the HTTP payload the system receives.
type CSPReport struct {
	Body CSPReportBody `json:"csp-report"`
}

// CSPReportBody contains the fields that are nested within the
// violation report.
type CSPReportBody struct {
	DocumentURI        string      `json:"document-uri"`
	Referrer           string      `json:"referrer"`
	BlockedURI         string      `json:"blocked-uri"`
	ViolatedDirective  string      `json:"violated-directive"`
	EffectiveDirective string      `json:"effective-directive"`
	OriginalPolicy     string      `json:"original-policy"`
	Disposition        string      `json:"disposition"`
	ScriptSample       string      `json:"script-sample"`
	StatusCode         interface{} `json:"status-code"`
}

const (
	// Default health check url.
	defaultHealthCheckPath = "/_healthcheck"
)

var (
	// Rev is set at build time and holds the revision that the package
	// was created at.
	Rev = "dev"

	// Shared defaults for the logger output. This ensures that we are
	// using the same keys for the `FieldKey` values across both formatters.
	logFieldMapDefaults = log.FieldMap{
		log.FieldKeyTime:  "timestamp",
		log.FieldKeyLevel: "level",
		log.FieldKeyMsg:   "message",
	}

	// Default URI Filter list.
	defaultIgnoredBlockedURIs = []string{
		"resource://",
		"chromenull://",
		"chrome-extension://",
		"safari-extension://",
		"mxjscall://",
		"webviewprogressproxy://",
		"res://",
		"mx://",
		"safari-resource://",
		"chromeinvoke://",
		"chromeinvokeimmediate://",
		"mbinit://",
		"opera://",
		"ms-appx://",
		"ms-appx-web://",
		"localhost",
		"127.0.0.1",
		"none://",
		"about:blank",
		"android-webview",
		"ms-browser-extension",
		"wvjbscheme://__wvjb_queue_message__",
		"nativebaiduhd://adblock",
		"bdvideo://error",
	}
)

func init() {
	log.SetOutput(os.Stdout)
	log.SetLevel(log.InfoLevel)
}

func trimEmptyAndComments(s []string) []string {
	var r []string
	for _, str := range s {
		if str == "" {
			continue
		}

		// ignore comments
		if strings.HasPrefix(str, "#") {
			continue
		}

		r = append(r, str)
	}
	return r
}

func main() {
	version := flag.Bool("version", false, "Display the version")
	debugFlag := flag.Bool("debug", false, "Output additional logging for debugging")
	outputFormat := flag.String("output-format", "text", "Define how the violation reports are formatted for output.\nDefaults to 'text'. Valid options are 'text' or 'json'")
	blockedURIFile := flag.String("filter-file", "", "Blocked URI Filter file")
	listenPort := flag.Int("port", 8080, "Port to listen on")
	healthCheckPath := flag.String("health-check-path", defaultHealthCheckPath, "Health checker path")
	truncateQueryStringFragment := flag.Bool("truncate-query-fragment", false, "Truncate query string and fragment from document-uri, referrer and blocked-uri before logging (to reduce chances of accidentally logging sensitive data)")

	flag.Parse()

	if *version {
		fmt.Printf("csp-collector (%s)\n", Rev)
		os.Exit(0)
	}

	if *debugFlag {
		log.SetLevel(log.DebugLevel)
	}

	if *outputFormat == "json" {
		log.SetFormatter(&log.JSONFormatter{
			FieldMap: logFieldMapDefaults,
		})
	} else {
		log.SetFormatter(&log.TextFormatter{
			FullTimestamp:          true,
			DisableLevelTruncation: true,
			QuoteEmptyFields:       true,
			DisableColors:          true,
			FieldMap:               logFieldMapDefaults,
		})
	}

	log.Debug("Starting up...")
	ignoredBlockedURIs := defaultIgnoredBlockedURIs
	if *blockedURIFile != "" {
		log.Debugf("Using Filter list from file at: %s\n", *blockedURIFile)

		content, err := ioutil.ReadFile(*blockedURIFile)
		if err != nil {
			log.Fatalf("Error reading Blocked File list: %s", *blockedURIFile)
		}
		ignoredBlockedURIs = trimEmptyAndComments(strings.Split(string(content), "\n"))
	} else {
		log.Debug("Using Filter list from internal list")
	}

	log.Debugf("Blocked URI List: %s", ignoredBlockedURIs)
	log.Debugf("Listening on TCP Port: %s", strconv.Itoa(*listenPort))

	http.HandleFunc(*healthCheckPath, func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}

		w.WriteHeader(http.StatusOK)
	})

	http.Handle("/", &violationReportHandler{
		blockedURIs:                 ignoredBlockedURIs,
		truncateQueryStringFragment: *truncateQueryStringFragment,
	})
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%s", strconv.Itoa(*listenPort)), nil))
}

type violationReportHandler struct {
	truncateQueryStringFragment bool
	blockedURIs                 []string
}

func (vrh *violationReportHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.WriteHeader(http.StatusMethodNotAllowed)
		log.WithFields(log.Fields{
			"http_method": r.Method,
		}).Debug("Received invalid HTTP method")
		return
	}

	decoder := json.NewDecoder(r.Body)
	var report CSPReport

	err := decoder.Decode(&report)
	if err != nil {
		w.WriteHeader(http.StatusUnprocessableEntity)
		log.Debugf("Unable to decode invalid JSON payload: %s", err)
		return
	}
	defer r.Body.Close()

	reportValidation := vrh.validateViolation(report)
	if reportValidation != nil {
		http.Error(w, reportValidation.Error(), http.StatusBadRequest)
		log.Debugf("Received invalid payload: %s", reportValidation.Error())
		return
	}

	metadatas, gotMetadata := r.URL.Query()["metadata"]
	var metadata string
	if gotMetadata {
		metadata = metadatas[0]
	}

	documentURI := report.Body.DocumentURI
	referrer := report.Body.Referrer
	blockedURI := report.Body.BlockedURI
	if vrh.truncateQueryStringFragment {
		documentURI = truncateQueryStringFragment(documentURI)
		referrer = truncateQueryStringFragment(referrer)
		blockedURI = truncateQueryStringFragment(blockedURI)
	}

	log.WithFields(log.Fields{
		"document_uri":        documentURI,
		"referrer":            referrer,
		"blocked_uri":         blockedURI,
		"violated_directive":  report.Body.ViolatedDirective,
		"effective_directive": report.Body.EffectiveDirective,
		"original_policy":     report.Body.OriginalPolicy,
		"disposition":         report.Body.Disposition,
		"script_sample":       report.Body.ScriptSample,
		"status_code":         report.Body.StatusCode,
		"metadata":            metadata,
		"path":                r.URL.Path,
	}).Info()
}

func (vrh *violationReportHandler) validateViolation(r CSPReport) error {
	for _, value := range vrh.blockedURIs {
		if strings.HasPrefix(r.Body.BlockedURI, value) {
			err := fmt.Errorf("blocked URI ('%s') is an invalid resource", value)
			return err
		}
	}

	if !strings.HasPrefix(r.Body.DocumentURI, "http") {
		return fmt.Errorf("document URI ('%s') is invalid", r.Body.DocumentURI)
	}

	return nil
}

func truncateQueryStringFragment(uri string) string {
	idx := strings.IndexAny(uri, "#?")
	if idx != -1 {
		return uri[:idx]
	}

	return uri
}
<commit_msg>Allow logging of (truncated) client IP<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/netip"
	"os"
	"strconv"
	"strings"

	log "github.com/sirupsen/logrus"
)

// CSPReport is the structure of the HTTP payload the system receives.
type CSPReport struct {
	Body CSPReportBody `json:"csp-report"`
}

// CSPReportBody contains the fields that are nested within the
// violation report.
type CSPReportBody struct {
	DocumentURI        string      `json:"document-uri"`
	Referrer           string      `json:"referrer"`
	BlockedURI         string      `json:"blocked-uri"`
	ViolatedDirective  string      `json:"violated-directive"`
	EffectiveDirective string      `json:"effective-directive"`
	OriginalPolicy     string      `json:"original-policy"`
	Disposition        string      `json:"disposition"`
	ScriptSample       string      `json:"script-sample"`
	StatusCode         interface{} `json:"status-code"`
}

const (
	// Default health check url.
	defaultHealthCheckPath = "/_healthcheck"
)

var (
	// Rev is set at build time and holds the revision that the package
	// was created at.
	Rev = "dev"

	// Shared defaults for the logger output. This ensures that we are
	// using the same keys for the `FieldKey` values across both formatters.
	logFieldMapDefaults = log.FieldMap{
		log.FieldKeyTime:  "timestamp",
		log.FieldKeyLevel: "level",
		log.FieldKeyMsg:   "message",
	}

	// Default URI Filter list.
	defaultIgnoredBlockedURIs = []string{
		"resource://",
		"chromenull://",
		"chrome-extension://",
		"safari-extension://",
		"mxjscall://",
		"webviewprogressproxy://",
		"res://",
		"mx://",
		"safari-resource://",
		"chromeinvoke://",
		"chromeinvokeimmediate://",
		"mbinit://",
		"opera://",
		"ms-appx://",
		"ms-appx-web://",
		"localhost",
		"127.0.0.1",
		"none://",
		"about:blank",
		"android-webview",
		"ms-browser-extension",
		"wvjbscheme://__wvjb_queue_message__",
		"nativebaiduhd://adblock",
		"bdvideo://error",
	}
)

func init() {
	log.SetOutput(os.Stdout)
	log.SetLevel(log.InfoLevel)
}

func trimEmptyAndComments(s []string) []string {
	var r []string
	for _, str := range s {
		if str == "" {
			continue
		}

		// ignore comments
		if strings.HasPrefix(str, "#") {
			continue
		}

		r = append(r, str)
	}
	return r
}

func main() {
	version := flag.Bool("version", false, "Display the version")
	debugFlag := flag.Bool("debug", false, "Output additional logging for debugging")
	outputFormat := flag.String("output-format", "text", "Define how the violation reports are formatted for output.\nDefaults to 'text'. Valid options are 'text' or 'json'")
	blockedURIFile := flag.String("filter-file", "", "Blocked URI Filter file")
	listenPort := flag.Int("port", 8080, "Port to listen on")
	healthCheckPath := flag.String("health-check-path", defaultHealthCheckPath, "Health checker path")
	truncateQueryStringFragment := flag.Bool("truncate-query-fragment", false, "Truncate query string and fragment from document-uri, referrer and blocked-uri before logging (to reduce chances of accidentally logging sensitive data)")

	logClientIP := flag.Bool("log-client-ip", false, "Log the reporting client IP address")
	logTruncatedClientIP := flag.Bool("log-truncated-client-ip", false, "Log the truncated client IP address (IPv4: /24, IPv6: /64")

	flag.Parse()

	if *version {
		fmt.Printf("csp-collector (%s)\n", Rev)
		os.Exit(0)
	}

	if *debugFlag {
		log.SetLevel(log.DebugLevel)
	}

	if *outputFormat == "json" {
		log.SetFormatter(&log.JSONFormatter{
			FieldMap: logFieldMapDefaults,
		})
	} else {
		log.SetFormatter(&log.TextFormatter{
			FullTimestamp:          true,
			DisableLevelTruncation: true,
			QuoteEmptyFields:       true,
			DisableColors:          true,
			FieldMap:               logFieldMapDefaults,
		})
	}

	log.Debug("Starting up...")
	ignoredBlockedURIs := defaultIgnoredBlockedURIs
	if *blockedURIFile != "" {
		log.Debugf("Using Filter list from file at: %s\n", *blockedURIFile)

		content, err := ioutil.ReadFile(*blockedURIFile)
		if err != nil {
			log.Fatalf("Error reading Blocked File list: %s", *blockedURIFile)
		}
		ignoredBlockedURIs = trimEmptyAndComments(strings.Split(string(content), "\n"))
	} else {
		log.Debug("Using Filter list from internal list")
	}

	log.Debugf("Blocked URI List: %s", ignoredBlockedURIs)
	log.Debugf("Listening on TCP Port: %s", strconv.Itoa(*listenPort))

	http.HandleFunc(*healthCheckPath, func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}

		w.WriteHeader(http.StatusOK)
	})

	http.Handle("/", &violationReportHandler{
		blockedURIs:                 ignoredBlockedURIs,
		truncateQueryStringFragment: *truncateQueryStringFragment,

		logClientIP:          *logClientIP,
		logTruncatedClientIP: *logTruncatedClientIP,
	})
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%s", strconv.Itoa(*listenPort)), nil))
}

type violationReportHandler struct {
	truncateQueryStringFragment bool
	blockedURIs                 []string

	logClientIP          bool
	logTruncatedClientIP bool
}

func (vrh *violationReportHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.WriteHeader(http.StatusMethodNotAllowed)
		log.WithFields(log.Fields{
			"http_method": r.Method,
		}).Debug("Received invalid HTTP method")
		return
	}

	decoder := json.NewDecoder(r.Body)
	var report CSPReport

	err := decoder.Decode(&report)
	if err != nil {
		w.WriteHeader(http.StatusUnprocessableEntity)
		log.Debugf("Unable to decode invalid JSON payload: %s", err)
		return
	}
	defer r.Body.Close()

	reportValidation := vrh.validateViolation(report)
	if reportValidation != nil {
		http.Error(w, reportValidation.Error(), http.StatusBadRequest)
		log.Debugf("Received invalid payload: %s", reportValidation.Error())
		return
	}

	metadatas, gotMetadata := r.URL.Query()["metadata"]
	var metadata string
	if gotMetadata {
		metadata = metadatas[0]
	}

	lf := log.Fields{
		"document_uri":       report.Body.DocumentURI,
		"referrer":            report.Body.Referrer,
		"blocked_uri":         report.Body.BlockedURI,
		"violated_directive":  report.Body.ViolatedDirective,
		"effective_directive": report.Body.EffectiveDirective,
		"original_policy":     report.Body.OriginalPolicy,
		"disposition":         report.Body.Disposition,
		"script_sample":       report.Body.ScriptSample,
		"status_code":         report.Body.StatusCode,
		"metadata":            metadata,
		"path":                r.URL.Path,
	}

	if vrh.truncateQueryStringFragment {
		lf["document_uri"] = truncateQueryStringFragment(report.Body.DocumentURI)
		lf["referrer"] = truncateQueryStringFragment(report.Body.Referrer)
		lf["blocked_uri"] = truncateQueryStringFragment(report.Body.BlockedURI)
	}

	if vrh.logClientIP {
		ip, err := getClientIP(r)
		if err != nil {
			log.Warnf("unable to parse client ip: %s", err)
		}
		lf["client_ip"] = ip.String()
	}

	if vrh.logTruncatedClientIP {
		ip, err := getClientIP(r)
		if err != nil {
			log.Warnf("unable to parse client ip: %s", err)
		}
		lf["client_ip"] = truncateClientIP(ip)
	}

	log.WithFields(lf).Info()
}

func (vrh *violationReportHandler) validateViolation(r CSPReport) error {
	for _, value := range vrh.blockedURIs {
		if strings.HasPrefix(r.Body.BlockedURI, value) {
			err := fmt.Errorf("blocked URI ('%s') is an invalid resource", value)
			return err
		}
	}

	if !strings.HasPrefix(r.Body.DocumentURI, "http") {
		return fmt.Errorf("document URI ('%s') is invalid", r.Body.DocumentURI)
	}

	return nil
}

func truncateQueryStringFragment(uri string) string {
	idx := strings.IndexAny(uri, "#?")
	if idx != -1 {
		return uri[:idx]
	}

	return uri
}

func truncateClientIP(a netip.Addr) string {
	// Ignoring the error is statically safe, as there are always enough bits.
	if a.Is4() {
		p, _ := a.Prefix(24)
		return p.String()
	}

	if a.Is6() {
		p, _ := a.Prefix(64)
		return p.String()
	}

	return "unknown-address"
}

func getClientIP(r *http.Request) (netip.Addr, error) {
	if s := r.Header.Get("X-Forwarded-For"); s != "" {
		return netip.ParseAddr(s)
	}

	addrp, err := netip.ParseAddrPort(r.RemoteAddr)
	return addrp.Addr(), err
}
<|endoftext|>


<commit_before>package worker

import (
	"context"
	"encoding/json"
	"fmt"
	"testing"
	"time"

	"github.com/pborman/uuid"
	workerctx "github.com/travis-ci/worker/context"
)

func TestAMQPLogWriterWrite(t *testing.T) {
	amqpConn, amqpChan := setupAMQPConn(t)
	defer amqpConn.Close()
	defer amqpChan.Close()

	uuid := uuid.NewRandom()
	ctx := workerctx.FromUUID(context.TODO(), uuid.String())

	logWriter, err := newAMQPLogWriter(ctx, amqpChan, 4, time.Hour, false)
	if err != nil {
		t.Fatal(err)
	}
	logWriter.SetMaxLogLength(1000)

	_, err = fmt.Fprintf(logWriter, "Hello, ")
	if err != nil {
		t.Error(err)
	}
	_, err = fmt.Fprintf(logWriter, "world!")
	if err != nil {
		t.Error(err)
	}

	// Close the log writer to force it to flush out the buffer
	err = logWriter.Close()
	if err != nil {
		t.Error(err)
	}

	delivery, ok, err := amqpChan.Get("reporting.jobs.logs", true)
	if err != nil {
		t.Error(err)
	}
	if !ok {
		t.Error("expected log message, but there was none")
	}

	var lp amqpLogPart

	err = json.Unmarshal(delivery.Body, &lp)
	if err != nil {
		t.Error(err)
	}

	expected := amqpLogPart{
		JobID:   4,
		Content: "Hello, world!",
		Number:  0,
		UUID:    uuid.String(),
		Final:   false,
	}

	if expected != lp {
		t.Errorf("log part is %#v, expected %#v", lp, expected)
	}
}

func TestAMQPLogWriterClose(t *testing.T) {
	amqpConn, amqpChan := setupAMQPConn(t)
	defer amqpConn.Close()
	defer amqpChan.Close()

	uuid := uuid.NewRandom()
	ctx := workerctx.FromUUID(context.TODO(), uuid.String())

	logWriter, err := newAMQPLogWriter(ctx, amqpChan, 4, time.Hour, false)
	if err != nil {
		t.Fatal(err)
	}
	logWriter.SetMaxLogLength(1000)

	// Close the log writer to force it to flush out the buffer
	err = logWriter.Close()
	if err != nil {
		t.Error(err)
	}

	delivery, ok, err := amqpChan.Get("reporting.jobs.logs", true)
	if err != nil {
		t.Error(err)
	}
	if !ok {
		t.Error("expected log message, but there was none")
	}

	var lp amqpLogPart

	err = json.Unmarshal(delivery.Body, &lp)
	if err != nil {
		t.Error(err)
	}

	expected := amqpLogPart{
		JobID:   4,
		Content: "",
		Number:  0,
		UUID:    uuid.String(),
		Final:   true,
	}

	if expected != lp {
		t.Errorf("log part is %#v, expected %#v", lp, expected)
	}
}

func TestAMQPMaxLogLength(t *testing.T) {
	amqpConn, amqpChan := setupAMQPConn(t)
	defer amqpConn.Close()
	defer amqpChan.Close()

	uuid := uuid.NewRandom()
	ctx := workerctx.FromUUID(context.TODO(), uuid.String())

	logWriter, err := newAMQPLogWriter(ctx, amqpChan, 4, time.Hour, false)
	if err != nil {
		t.Fatal(err)
	}
	logWriter.SetMaxLogLength(4)

	_, err = fmt.Fprintf(logWriter, "1234")
	if err != nil {
		t.Error(err)
	}
	_, err = fmt.Fprintf(logWriter, "5")
	if err == nil {
		t.Error("expected error, but got nil")
	}
}
<commit_msg>add missing queuedAt arg to tests<commit_after>package worker

import (
	"context"
	"encoding/json"
	"fmt"
	"testing"
	"time"

	"github.com/pborman/uuid"
	workerctx "github.com/travis-ci/worker/context"
)

func TestAMQPLogWriterWrite(t *testing.T) {
	amqpConn, amqpChan := setupAMQPConn(t)
	defer amqpConn.Close()
	defer amqpChan.Close()

	uuid := uuid.NewRandom()
	queuedAt := time.Now()
	ctx := workerctx.FromUUID(context.TODO(), uuid.String())

	logWriter, err := newAMQPLogWriter(ctx, amqpChan, 4, &queuedAt, time.Hour, false)
	if err != nil {
		t.Fatal(err)
	}
	logWriter.SetMaxLogLength(1000)

	_, err = fmt.Fprintf(logWriter, "Hello, ")
	if err != nil {
		t.Error(err)
	}
	_, err = fmt.Fprintf(logWriter, "world!")
	if err != nil {
		t.Error(err)
	}

	// Close the log writer to force it to flush out the buffer
	err = logWriter.Close()
	if err != nil {
		t.Error(err)
	}

	delivery, ok, err := amqpChan.Get("reporting.jobs.logs", true)
	if err != nil {
		t.Error(err)
	}
	if !ok {
		t.Error("expected log message, but there was none")
	}

	var lp amqpLogPart

	err = json.Unmarshal(delivery.Body, &lp)
	if err != nil {
		t.Error(err)
	}

	expected := amqpLogPart{
		JobID:   4,
		Content: "Hello, world!",
		Number:  0,
		UUID:    uuid.String(),
		Final:   false,
	}

	if expected != lp {
		t.Errorf("log part is %#v, expected %#v", lp, expected)
	}
}

func TestAMQPLogWriterClose(t *testing.T) {
	amqpConn, amqpChan := setupAMQPConn(t)
	defer amqpConn.Close()
	defer amqpChan.Close()

	uuid := uuid.NewRandom()
	queuedAt := time.Now()
	ctx := workerctx.FromUUID(context.TODO(), uuid.String())

	logWriter, err := newAMQPLogWriter(ctx, amqpChan, 4, &queuedAt, time.Hour, false)
	if err != nil {
		t.Fatal(err)
	}
	logWriter.SetMaxLogLength(1000)

	// Close the log writer to force it to flush out the buffer
	err = logWriter.Close()
	if err != nil {
		t.Error(err)
	}

	delivery, ok, err := amqpChan.Get("reporting.jobs.logs", true)
	if err != nil {
		t.Error(err)
	}
	if !ok {
		t.Error("expected log message, but there was none")
	}

	var lp amqpLogPart

	err = json.Unmarshal(delivery.Body, &lp)
	if err != nil {
		t.Error(err)
	}

	expected := amqpLogPart{
		JobID:   4,
		Content: "",
		Number:  0,
		UUID:    uuid.String(),
		Final:   true,
	}

	if expected != lp {
		t.Errorf("log part is %#v, expected %#v", lp, expected)
	}
}

func TestAMQPMaxLogLength(t *testing.T) {
	amqpConn, amqpChan := setupAMQPConn(t)
	defer amqpConn.Close()
	defer amqpChan.Close()

	uuid := uuid.NewRandom()
	queuedAt := time.Now()
	ctx := workerctx.FromUUID(context.TODO(), uuid.String())

	logWriter, err := newAMQPLogWriter(ctx, amqpChan, 4, &queuedAt, time.Hour, false)
	if err != nil {
		t.Fatal(err)
	}
	logWriter.SetMaxLogLength(4)

	_, err = fmt.Fprintf(logWriter, "1234")
	if err != nil {
		t.Error(err)
	}
	_, err = fmt.Fprintf(logWriter, "5")
	if err == nil {
		t.Error("expected error, but got nil")
	}
}
<|endoftext|>


<commit_before>// An experimental FTP server framework. By providing a simple driver class that
// responds to a handful of methods you can have a complete FTP server.
//
// Some sample use cases include persisting data to an Amazon S3 bucket, a
// relational database, redis or memory.
//
// There is a sample in-memory driver available - see the documentation for the
// graval-mem package or the graval READEME for more details.
package graval

import (
	"crypto/tls"
	"net"
	"strconv"
	"strings"
)

type CryptoConfig struct {
	Implicit  bool
	TlsConfig *tls.Config
}

// serverOpts contains parameters for graval.NewFTPServer()
type FTPServerOpts struct {
	// Server name will be used for welcome message
	ServerName string

	// The factory that will be used to create a new FTPDriver instance for
	// each client connection. This is a mandatory option.
	Factory FTPDriverFactory

	// The hostname that the FTP server should listen on. Optional, defaults to
	// "::", which means all hostnames on ipv4 and ipv6.
	Hostname string

	// The port that the FTP should listen on. Optional, defaults to 3000. In
	// a production environment you will probably want to change this to 21.
	Port int

	// Options for passive data connections
	PassiveOpts *PassiveOpts

	// Options for FTPS and FTPES
	CryptoConfig *CryptoConfig
}

// FTPServer is the root of your FTP application. You should instantiate one
// of these and call ListenAndServe() to start accepting client connections.
//
// Always use the NewFTPServer() method to create a new FTPServer.
type FTPServer struct {
	serverName    string
	listenTo      string
	driverFactory FTPDriverFactory
	logger        *ftpLogger
	passiveOpts   *PassiveOpts
	cryptoConfig  *CryptoConfig
}

// serverOptsWithDefaults copies an FTPServerOpts struct into a new struct,
// then adds any default values that are missing and returns the new data.
func serverOptsWithDefaults(opts *FTPServerOpts) *FTPServerOpts {
	var newOpts FTPServerOpts

	if opts == nil {
		opts = &FTPServerOpts{}
	}

	if opts.ServerName == "" {
		newOpts.ServerName = "Go FTP Server"
	} else {
		newOpts.ServerName = opts.ServerName
	}

	if opts.Hostname == "" {
		newOpts.Hostname = "::"
	} else {
		newOpts.Hostname = opts.Hostname
	}

	if opts.Port == 0 {
		newOpts.Port = 3000
	} else {
		newOpts.Port = opts.Port
	}

	newOpts.Factory = opts.Factory

	if opts.PassiveOpts == nil {
		newOpts.PassiveOpts = &PassiveOpts{}
	} else {
		newOpts.PassiveOpts = opts.PassiveOpts
	}

	newOpts.CryptoConfig = opts.CryptoConfig

	return &newOpts
}

// NewFTPServer initialises a new FTP server. Configuration options are provided
// via an instance of FTPServerOpts. Calling this function in your code will
// probably look something like this:
//
//     factory := &MyDriverFactory{}
//     server  := graval.NewFTPServer(&graval.FTPServerOpts{ Factory: factory })
//
// or:
//
//     factory := &MyDriverFactory{}
//     opts    := &graval.FTPServerOpts{
//       Factory: factory,
//       Port: 2000,
//       Hostname: "127.0.0.1",
//     }
//     server  := graval.NewFTPServer(opts)
//
func NewFTPServer(opts *FTPServerOpts) *FTPServer {
	opts = serverOptsWithDefaults(opts)
	s := new(FTPServer)
	s.listenTo = buildTcpString(opts.Hostname, opts.Port)
	s.serverName = opts.ServerName
	s.driverFactory = opts.Factory
	s.logger = newFtpLogger("")
	s.passiveOpts = opts.PassiveOpts
	s.cryptoConfig = opts.CryptoConfig
	return s
}

// ListenAndServe asks a new FTPServer to begin accepting client connections. It
// accepts no arguments - all configuration is provided via the NewFTPServer
// function.
//
// If the server fails to start for any reason, an error will be returned. Common
// errors are trying to bind to a privileged port or something else is already
// listening on the same port.
//
func (ftpServer *FTPServer) ListenAndServe() error {
	laddr, err := net.ResolveTCPAddr("tcp", ftpServer.listenTo)
	if err != nil {
		return err
	}
	listener, err := net.ListenTCP("tcp", laddr)
	if err != nil {
		return err
	}
	for {
		tcpConn, err := listener.AcceptTCP()
		if err != nil {
			ftpServer.logger.Print("listening error")
			break
		}
		driver, err := ftpServer.driverFactory.NewDriver()
		if err != nil {
			ftpServer.logger.Print("Error creating driver, aborting client connection")
		} else {
			ftpConn := newftpConn(tcpConn, driver, ftpServer.serverName, ftpServer.passiveOpts, ftpServer.cryptoConfig)
			go ftpConn.Serve()
		}
	}
	return nil
}

func buildTcpString(hostname string, port int) (result string) {
	if strings.Contains(hostname, ":") {
		// ipv6
		if port == 0 {
			result = "[" + hostname + "]"
		} else {
			result = "[" + hostname + "]:" + strconv.Itoa(port)
		}
	} else {
		// ipv4
		if port == 0 {
			result = hostname
		} else {
			result = hostname + ":" + strconv.Itoa(port)
		}
	}
	return
}
<commit_msg>FTP server Close<commit_after>// An experimental FTP server framework. By providing a simple driver class that
// responds to a handful of methods you can have a complete FTP server.
//
// Some sample use cases include persisting data to an Amazon S3 bucket, a
// relational database, redis or memory.
//
// There is a sample in-memory driver available - see the documentation for the
// graval-mem package or the graval READEME for more details.
package graval

import (
	"crypto/tls"
	"net"
	"strconv"
	"strings"
)

type CryptoConfig struct {
	Implicit  bool
	TlsConfig *tls.Config
}

// serverOpts contains parameters for graval.NewFTPServer()
type FTPServerOpts struct {
	// Server name will be used for welcome message
	ServerName string

	// The factory that will be used to create a new FTPDriver instance for
	// each client connection. This is a mandatory option.
	Factory FTPDriverFactory

	// The hostname that the FTP server should listen on. Optional, defaults to
	// "::", which means all hostnames on ipv4 and ipv6.
	Hostname string

	// The port that the FTP should listen on. Optional, defaults to 3000. In
	// a production environment you will probably want to change this to 21.
	Port int

	// Options for passive data connections
	PassiveOpts *PassiveOpts

	// Options for FTPS and FTPES
	CryptoConfig *CryptoConfig
}

// FTPServer is the root of your FTP application. You should instantiate one
// of these and call ListenAndServe() to start accepting client connections.
//
// Always use the NewFTPServer() method to create a new FTPServer.
type FTPServer struct {
	serverName    string
	listenTo      string
	listener      net.Listener
	closed        bool
	driverFactory FTPDriverFactory
	logger        *ftpLogger
	passiveOpts   *PassiveOpts
	cryptoConfig  *CryptoConfig
}

// serverOptsWithDefaults copies an FTPServerOpts struct into a new struct,
// then adds any default values that are missing and returns the new data.
func serverOptsWithDefaults(opts *FTPServerOpts) *FTPServerOpts {
	var newOpts FTPServerOpts

	if opts == nil {
		opts = &FTPServerOpts{}
	}

	if opts.ServerName == "" {
		newOpts.ServerName = "Go FTP Server"
	} else {
		newOpts.ServerName = opts.ServerName
	}

	if opts.Hostname == "" {
		newOpts.Hostname = "::"
	} else {
		newOpts.Hostname = opts.Hostname
	}

	if opts.Port == 0 {
		newOpts.Port = 3000
	} else {
		newOpts.Port = opts.Port
	}

	newOpts.Factory = opts.Factory

	if opts.PassiveOpts == nil {
		newOpts.PassiveOpts = &PassiveOpts{}
	} else {
		newOpts.PassiveOpts = opts.PassiveOpts
	}

	newOpts.CryptoConfig = opts.CryptoConfig

	return &newOpts
}

// NewFTPServer initialises a new FTP server. Configuration options are provided
// via an instance of FTPServerOpts. Calling this function in your code will
// probably look something like this:
//
//     factory := &MyDriverFactory{}
//     server  := graval.NewFTPServer(&graval.FTPServerOpts{ Factory: factory })
//
// or:
//
//     factory := &MyDriverFactory{}
//     opts    := &graval.FTPServerOpts{
//       Factory: factory,
//       Port: 2000,
//       Hostname: "127.0.0.1",
//     }
//     server  := graval.NewFTPServer(opts)
//
func NewFTPServer(opts *FTPServerOpts) *FTPServer {
	opts = serverOptsWithDefaults(opts)
	s := new(FTPServer)
	s.listenTo = buildTcpString(opts.Hostname, opts.Port)
	s.serverName = opts.ServerName
	s.driverFactory = opts.Factory
	s.logger = newFtpLogger("")
	s.passiveOpts = opts.PassiveOpts
	s.cryptoConfig = opts.CryptoConfig
	return s
}

// ListenAndServe asks a new FTPServer to begin accepting client connections. It
// accepts no arguments - all configuration is provided via the NewFTPServer
// function.
//
// If the server fails to start for any reason, an error will be returned. Common
// errors are trying to bind to a privileged port or something else is already
// listening on the same port.
//
func (ftpServer *FTPServer) ListenAndServe() error {
	laddr, err := net.ResolveTCPAddr("tcp", ftpServer.listenTo)
	if err != nil {
		return err
	}
	listener, err := net.ListenTCP("tcp", laddr)
	if err != nil {
		return err
	}
	ftpServer.listener = listener
	for {
		tcpConn, err := listener.AcceptTCP()

		if ftpServer.closed {
			return nil
		}

		if err != nil {
			ftpServer.logger.Print("listening error")
			break
		}
		driver, err := ftpServer.driverFactory.NewDriver()
		if err != nil {
			ftpServer.logger.Print("Error creating driver, aborting client connection")
		} else {
			ftpConn := newftpConn(tcpConn, driver, ftpServer.serverName, ftpServer.passiveOpts, ftpServer.cryptoConfig)
			go ftpConn.Serve()
		}
	}
	return nil
}

func (ftpServer *FTPServer) Close() error {
	ftpServer.closed = true
	return ftpServer.listener.Close()
}

func buildTcpString(hostname string, port int) (result string) {
	if strings.Contains(hostname, ":") {
		// ipv6
		if port == 0 {
			result = "[" + hostname + "]"
		} else {
			result = "[" + hostname + "]:" + strconv.Itoa(port)
		}
	} else {
		// ipv4
		if port == 0 {
			result = hostname
		} else {
			result = hostname + ":" + strconv.Itoa(port)
		}
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
)

type app_state struct {
	Wait    sync.WaitGroup
	Running bool
}

var (
	InfoLogger  *log.Logger
	ErrorLogger *log.Logger
	DebugLogger *log.Logger
)

func InitLogging(WithDebug bool) {
	InfoLogger = log.New(os.Stdout, "", log.Ldate|log.Ltime|log.Lshortfile)
	ErrorLogger = log.New(os.Stderr, "", log.Ldate|log.Ltime|log.Lshortfile)

	var DebugWriter io.Writer
	if WithDebug {
		DebugWriter = os.Stdout
	} else {
		DebugWriter = ioutil.Discard
	}
	DebugLogger = log.New(DebugWriter, "", log.Ldate|log.Ltime|log.Lshortfile)

}

func KillSignal(Running *bool) {
	SignalChannel := make(chan os.Signal, 1)
	signal.Notify(SignalChannel, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		<-SignalChannel
		InfoLogger.Println("Got SIGTERM, finishing work gracefully.")
		*Running = false
		signal.Reset() // Next ctrl+c will effect in ungraceful stop
	}()

}

func MainLoop(AppState *app_state, AppConfig app_config) {
	AppState.Running = true

	AppState.Wait.Add(1)
	go sFlowListener(AppState, AppConfig.SFlowConfig)

	AppState.Wait.Wait()
}

func main() {
	var (
		configFile string
		withDebug  bool
		AppState   app_state
		AppConfig  app_config
		err        error
	)

	flag.StringVar(&configFile, "c", "/etc/ddos_detector.toml", "Path to configuration file.")
	flag.BoolVar(&withDebug, "v", false, "Be verbose, show debugging output.")
	flag.Parse()

	InitLogging(withDebug)
	InfoLogger.Println("Starting DDoS Detector.")

	AppConfig, err = ReadConfig(configFile)
	if err != nil {
		os.Exit(1)
	}

	KillSignal(&AppState.Running)
	MainLoop(&AppState, AppConfig)

	InfoLogger.Println("DDoS Detector finished.")
}
<commit_msg>Stop endless copying of structures<commit_after>package main

import (
	"flag"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
)

type app_state struct {
	Wait    sync.WaitGroup
	Running bool
}

var (
	InfoLogger  *log.Logger
	ErrorLogger *log.Logger
	DebugLogger *log.Logger
)

func InitLogging(WithDebug bool) {
	InfoLogger = log.New(os.Stdout, "", log.Ldate|log.Ltime|log.Lshortfile)
	ErrorLogger = log.New(os.Stderr, "", log.Ldate|log.Ltime|log.Lshortfile)

	var DebugWriter io.Writer
	if WithDebug {
		DebugWriter = os.Stdout
	} else {
		DebugWriter = ioutil.Discard
	}
	DebugLogger = log.New(DebugWriter, "", log.Ldate|log.Ltime|log.Lshortfile)

}

func KillSignal(Running *bool) {
	SignalChannel := make(chan os.Signal, 1)
	signal.Notify(SignalChannel, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		<-SignalChannel
		InfoLogger.Println("Got SIGTERM, finishing work gracefully.")
		*Running = false
		signal.Reset() // Next ctrl+c will effect in ungraceful stop
	}()

}

func MainLoop(AppState *app_state, AppConfig *app_config) {
	AppState.Running = true

	AppState.Wait.Add(1)
	go sFlowListener(AppState, AppConfig.SFlowConfig)

	AppState.Wait.Wait()
}

func main() {
	var (
		configFile string
		withDebug  bool
		AppState   app_state
		AppConfig  app_config
		err        error
	)

	flag.StringVar(&configFile, "c", "/etc/ddos_detector.toml", "Path to configuration file.")
	flag.BoolVar(&withDebug, "v", false, "Be verbose, show debugging output.")
	flag.Parse()

	InitLogging(withDebug)
	InfoLogger.Println("Starting DDoS Detector.")

	AppConfig, err = ReadConfig(configFile)
	if err != nil {
		os.Exit(1)
	}

	KillSignal(&AppState.Running)
	MainLoop(&AppState, &AppConfig)

	InfoLogger.Println("DDoS Detector finished.")
}
<|endoftext|>


<commit_before>package binny

import (
	"encoding"
	"encoding/gob"
	"fmt"
	"log"
	"reflect"
	"sync"
)

type DecoderTypeError struct {
	Expected string
	Actual   Type
}

func (dte DecoderTypeError) Error() string {
	return "expected " + dte.Expected + ", got " + dte.Actual.String()
}

type decoderFunc func(dec *Decoder, v reflect.Value) error

var (
	decCache = struct {
		sync.RWMutex
		m map[reflect.Type]decoderFunc
	}{m: map[reflect.Type]decoderFunc{}}
)

func typeDecoder(t reflect.Type) (fn decoderFunc) {
	decCache.RLock()
	fn = decCache.m[t]
	decCache.RUnlock()
	if fn != nil {
		return
	}

	decCache.Lock()
	decCache.m[t] = nopeDec
	decCache.Unlock()

	fn = newTypeDecoder(t)
	decCache.Lock()
	decCache.m[t] = fn
	if k := t.Kind(); k == reflect.Struct || k == reflect.Ptr && t.Elem().Kind() == reflect.Struct {
		updateCachedFields(false)
	}
	decCache.Unlock()
	return
}

func nopeDec(*Decoder, reflect.Value) error { panic("Fly, you fools!") }

func newTypeDecoder(t reflect.Type) decoderFunc {
	log.Println(t)
	k := t.Kind()
	if t.Implements(unmarshalerType) {
		if k == reflect.Ptr {
			return newPtrDecoder(unmarshalerDecoder, false)
		}
		return unmarshalerDecoder
	}

	if t.Implements(binaryUnmarshalerType) {
		if k == reflect.Ptr {
			return newPtrDecoder(binaryUnmarshalerDecoder, false)
		}
		return binaryUnmarshalerDecoder
	}
	if t.Implements(gobDecoderType) {
		if k == reflect.Ptr {
			return newPtrDecoder(gobDecoder, false)
		}
		return gobDecoder
	}

	if k != reflect.Ptr {
		t := reflect.PtrTo(t)
		if t.Implements(unmarshalerType) {
			return addrDecoder(unmarshalerDecoder)
		}
		if t.Implements(binaryUnmarshalerType) {
			return addrDecoder(binaryUnmarshalerDecoder)
		}
		if t.Implements(gobDecoderType) {
			return addrDecoder(gobDecoder)
		}
	}
	switch k {
	case reflect.Bool:
		return boolDecoder
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return intDecoder
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
		return uintDecoder
	case reflect.Float32:
		return float32Decoder
	case reflect.Float64:
		return float64Decoder
	case reflect.Complex64:
		return complex64Decoder
	case reflect.Complex128:
		return complex128Decoder
	case reflect.String:
		return stringDecoder
	case reflect.Map:
		if !isNative(t.Key().Kind(), true) { // eventually will support pointer keys and structs, but not gonna happen until it's needed
			return invalidDecoder
		}
		return newMapDecoder(t)
	case reflect.Slice, reflect.Array:
		return newSliceDecoder(t.Elem())
	case reflect.Struct:
		return newStructDecoder(t)
	case reflect.Ptr:
		return newPtrDecoder(typeDecoder(t.Elem()), true)
		// case reflect.Interface:
		// 	return ifaceEncoder
	}
	return invalidDecoder
}

func unmarshalerDecoder(d *Decoder, v reflect.Value) error {
	return v.Interface().(Unmarshaler).UnmarshalBinny(d)
}

func binaryUnmarshalerDecoder(d *Decoder, v reflect.Value) error {
	return d.ReadBinary(v.Interface().(encoding.BinaryUnmarshaler))
}

func gobDecoder(d *Decoder, v reflect.Value) error {
	return d.ReadGob(v.Interface().(gob.GobDecoder))
}

func boolDecoder(d *Decoder, v reflect.Value) error {
	b, err := d.ReadBool()
	v.SetBool(b)
	return err
}

func intDecoder(d *Decoder, v reflect.Value) error {
	i, _, err := d.ReadInt()
	v.SetInt(i)
	return err
}

func uintDecoder(d *Decoder, v reflect.Value) error {
	i, _, err := d.ReadUint()
	v.SetUint(i)
	return err
}

func float32Decoder(d *Decoder, v reflect.Value) error {
	f, err := d.ReadFloat32()
	v.SetFloat(float64(f))
	return err
}

func float64Decoder(d *Decoder, v reflect.Value) error {
	f, err := d.ReadFloat64()
	v.SetFloat(f)
	return err
}

func complex64Decoder(d *Decoder, v reflect.Value) error {
	c, err := d.ReadComplex64()
	v.SetComplex(complex128(c))
	return err
}

func complex128Decoder(d *Decoder, v reflect.Value) error {
	c, err := d.ReadComplex128()
	v.SetComplex(c)
	return err
}

func stringDecoder(d *Decoder, v reflect.Value) error {
	s, err := d.ReadString()
	v.SetString(s)
	return err
}

func bytesDecoder(d *Decoder, v reflect.Value) error {
	b, err := d.ReadBytes()
	v.SetBytes(b)
	return err
}

func invalidDecoder(d *Decoder, v reflect.Value) error {
	return fmt.Errorf("%v is not supported", v.Type().String())
}

type sliceDecoder struct {
	t reflect.Type
}

func (sd sliceDecoder) decode(d *Decoder, v reflect.Value) error {
	if err := d.expectType(Slice); err != nil {
		return err
	}

	ln, _, err := d.ReadUint()

	if err != nil {
		return err
	}

	if v.Kind() == reflect.Slice {
		if ln := int(ln); v.Cap() < ln {
			v.Set(reflect.MakeSlice(v.Type(), ln, ln))
		} else {
			v.SetLen(ln)
			v.SetCap(ln)
		}
	}
	dec := typeDecoder(sd.t)
	for i := 0; i < int(ln); i++ {
		// this is a bug
		if d.peekType() == Nil {
			d.readType()
			continue
		}
		if err = dec(d, v.Index(i)); err != nil {
			return err
		}
	}

	return d.expectType(EOV)
}

func newSliceDecoder(t reflect.Type) decoderFunc {
	if t.Kind() == reflect.Uint8 {
		return bytesDecoder
	}
	typeDecoder(t)
	d := sliceDecoder{t: t}
	return d.decode
}

type structDecoder struct {
	fields map[string]*field
}

func (sd structDecoder) decode(d *Decoder, v reflect.Value) error {
	if err := d.expectType(Struct); err != nil {
		if err, ok := err.(DecoderTypeError); ok && err.Actual == Nil || err.Actual == EmptyStruct {
			return nil
		}
		return err
	}
	for {
		n, err := d.ReadString()
		if err != nil {
			if err, ok := err.(DecoderTypeError); ok && err.Actual == EOV {
				return nil
			}
			return err
		}
		if f, ok := sd.fields[n]; ok {
			fld := fieldByIndex(v, f.index, true)
			if err := f.dec(d, fld); err != nil {
				return err
			}
		}
	}
}

func newStructDecoder(t reflect.Type) decoderFunc {
	var (
		flds = cachedTypeFields(t)
		sd   = structDecoder{fields: make(map[string]*field, len(flds))}
	)
	for i := range flds {
		f := &flds[i] // need a pointer so when we later update fields .dec it'd get picked up.
		sd.fields[f.name] = f
	}
	return sd.decode
}

type mapDecoder struct {
	kt, vt reflect.Type
}

func (md mapDecoder) decode(d *Decoder, v reflect.Value) error {
	if err := d.expectType(Map); err != nil {
		return err
	}
	ln, _, err := d.ReadUint()

	if err != nil {
		return err
	}

	t := v.Type()
	if v.IsNil() {
		v.Set(reflect.MakeMap(t))
	}
	kdec, vdec := typeDecoder(md.kt), typeDecoder(md.vt)
	for i, kt, vt := 0, t.Key(), t.Elem(); i < int(ln); i++ {
		key := reflect.New(kt).Elem()
		if err = kdec(d, key); err != nil {
			return err
		}
		if d.peekType() == Nil {
			v.SetMapIndex(key, reflect.Zero(vt))
			d.readType()
			continue
		}
		val := reflect.New(vt).Elem()
		if err = vdec(d, val); err != nil {
			return err
		}
		v.SetMapIndex(key, val)
	}
	return d.expectType(EOV)
}

func newMapDecoder(t reflect.Type) decoderFunc {
	md := mapDecoder{t.Key(), t.Elem()}
	return md.decode
}

type ptrDecoder struct {
	dec decoderFunc
}

func (pd ptrDecoder) decodeElem(d *Decoder, v reflect.Value) error {
	if v.IsNil() {
		v.Set(reflect.New(v.Type().Elem()))
	}
	return pd.dec(d, v.Elem())
}

func (pd ptrDecoder) decode(d *Decoder, v reflect.Value) error {
	if !v.CanAddr() && v.IsNil() {
		v.Set(reflect.New(v.Type().Elem()))
	}
	return pd.dec(d, v)
}

func newPtrDecoder(fn decoderFunc, elem bool) decoderFunc {
	pd := ptrDecoder{fn}
	if elem {
		return pd.decodeElem
	}
	return pd.decode
}

func addrDecoder(fn decoderFunc) decoderFunc {
	return func(d *Decoder, v reflect.Value) error {
		return fn(d, v.Addr())
	}
}
<commit_msg>fix addr decoding<commit_after>package binny

import (
	"encoding"
	"encoding/gob"
	"fmt"
	"reflect"
	"sync"
)

type DecoderTypeError struct {
	Expected string
	Actual   Type
}

func (dte DecoderTypeError) Error() string {
	return "expected " + dte.Expected + ", got " + dte.Actual.String()
}

type decoderFunc func(dec *Decoder, v reflect.Value) error

var (
	decCache = struct {
		sync.RWMutex
		m map[reflect.Type]decoderFunc
	}{m: map[reflect.Type]decoderFunc{}}
)

func typeDecoder(t reflect.Type) (fn decoderFunc) {
	decCache.RLock()
	fn = decCache.m[t]
	decCache.RUnlock()
	if fn != nil {
		return
	}

	decCache.Lock()
	decCache.m[t] = nopeDec
	decCache.Unlock()

	fn = newTypeDecoder(t)
	decCache.Lock()
	decCache.m[t] = fn
	if k := t.Kind(); k == reflect.Struct || k == reflect.Ptr && t.Elem().Kind() == reflect.Struct {
		updateCachedFields(false)
	}
	decCache.Unlock()
	return
}

func nopeDec(*Decoder, reflect.Value) error { panic("Fly, you fools!") }

func newTypeDecoder(t reflect.Type) decoderFunc {
	k := t.Kind()
	if t.Implements(unmarshalerType) {
		if k == reflect.Ptr {
			return newPtrDecoder(unmarshalerDecoder, false)
		}
		return unmarshalerDecoder
	}

	if t.Implements(binaryUnmarshalerType) {
		if k == reflect.Ptr {
			return newPtrDecoder(binaryUnmarshalerDecoder, false)
		}
		return binaryUnmarshalerDecoder
	}
	if t.Implements(gobDecoderType) {
		if k == reflect.Ptr {
			return newPtrDecoder(gobDecoder, false)
		}
		return gobDecoder
	}

	if k != reflect.Ptr {
		t := reflect.PtrTo(t)
		if t.Implements(unmarshalerType) {
			return addrDecoder(unmarshalerDecoder)
		}
		if t.Implements(binaryUnmarshalerType) {
			return addrDecoder(binaryUnmarshalerDecoder)
		}
		if t.Implements(gobDecoderType) {
			return addrDecoder(gobDecoder)
		}
	}
	switch k {
	case reflect.Bool:
		return boolDecoder
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return intDecoder
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
		return uintDecoder
	case reflect.Float32:
		return float32Decoder
	case reflect.Float64:
		return float64Decoder
	case reflect.Complex64:
		return complex64Decoder
	case reflect.Complex128:
		return complex128Decoder
	case reflect.String:
		return stringDecoder
	case reflect.Map:
		if !isNative(t.Key().Kind(), true) { // eventually will support pointer keys and structs, but not gonna happen until it's needed
			return invalidDecoder
		}
		return newMapDecoder(t)
	case reflect.Slice, reflect.Array:
		return newSliceDecoder(t.Elem())
	case reflect.Struct:
		return newStructDecoder(t)
	case reflect.Ptr:
		return newPtrDecoder(typeDecoder(t.Elem()), true)
		// case reflect.Interface:
		// 	return ifaceEncoder
	}
	return invalidDecoder
}

func unmarshalerDecoder(d *Decoder, v reflect.Value) error {
	return v.Interface().(Unmarshaler).UnmarshalBinny(d)
}

func binaryUnmarshalerDecoder(d *Decoder, v reflect.Value) error {
	return d.ReadBinary(v.Interface().(encoding.BinaryUnmarshaler))
}

func gobDecoder(d *Decoder, v reflect.Value) error {
	return d.ReadGob(v.Interface().(gob.GobDecoder))
}

func boolDecoder(d *Decoder, v reflect.Value) error {
	b, err := d.ReadBool()
	v.SetBool(b)
	return err
}

func intDecoder(d *Decoder, v reflect.Value) error {
	i, _, err := d.ReadInt()
	v.SetInt(i)
	return err
}

func uintDecoder(d *Decoder, v reflect.Value) error {
	i, _, err := d.ReadUint()
	v.SetUint(i)
	return err
}

func float32Decoder(d *Decoder, v reflect.Value) error {
	f, err := d.ReadFloat32()
	v.SetFloat(float64(f))
	return err
}

func float64Decoder(d *Decoder, v reflect.Value) error {
	f, err := d.ReadFloat64()
	v.SetFloat(f)
	return err
}

func complex64Decoder(d *Decoder, v reflect.Value) error {
	c, err := d.ReadComplex64()
	v.SetComplex(complex128(c))
	return err
}

func complex128Decoder(d *Decoder, v reflect.Value) error {
	c, err := d.ReadComplex128()
	v.SetComplex(c)
	return err
}

func stringDecoder(d *Decoder, v reflect.Value) error {
	s, err := d.ReadString()
	v.SetString(s)
	return err
}

func bytesDecoder(d *Decoder, v reflect.Value) error {
	b, err := d.ReadBytes()
	v.SetBytes(b)
	return err
}

func invalidDecoder(d *Decoder, v reflect.Value) error {
	return fmt.Errorf("%v is not supported", v.Type().String())
}

type sliceDecoder struct {
	t reflect.Type
}

func (sd sliceDecoder) decode(d *Decoder, v reflect.Value) error {
	if err := d.expectType(Slice); err != nil {
		return err
	}

	ln, _, err := d.ReadUint()

	if err != nil {
		return err
	}

	if v.Kind() == reflect.Slice {
		if ln := int(ln); v.Cap() < ln {
			v.Set(reflect.MakeSlice(v.Type(), ln, ln))
		} else {
			v.SetLen(ln)
			v.SetCap(ln)
		}
	}
	dec := typeDecoder(sd.t)
	for i := 0; i < int(ln); i++ {
		// this is a bug
		if d.peekType() == Nil {
			d.readType()
			continue
		}
		if err = dec(d, v.Index(i)); err != nil {
			return err
		}
	}

	return d.expectType(EOV)
}

func newSliceDecoder(t reflect.Type) decoderFunc {
	if t.Kind() == reflect.Uint8 {
		return bytesDecoder
	}
	typeDecoder(t)
	d := sliceDecoder{t: t}
	return d.decode
}

type structDecoder struct {
	fields map[string]*field
}

func (sd structDecoder) decode(d *Decoder, v reflect.Value) error {
	if err := d.expectType(Struct); err != nil {
		if err, ok := err.(DecoderTypeError); ok && err.Actual == Nil || err.Actual == EmptyStruct {
			return nil
		}
		return err
	}
	for {
		n, err := d.ReadString()
		if err != nil {
			if err, ok := err.(DecoderTypeError); ok && err.Actual == EOV {
				return nil
			}
			return err
		}
		if f, ok := sd.fields[n]; ok {
			fld := fieldByIndex(v, f.index, true)
			if err := f.dec(d, fld); err != nil {
				return err
			}
		}
	}
}

func newStructDecoder(t reflect.Type) decoderFunc {
	var (
		flds = cachedTypeFields(t)
		sd   = structDecoder{fields: make(map[string]*field, len(flds))}
	)
	for i := range flds {
		f := &flds[i] // need a pointer so when we later update fields .dec it'd get picked up.
		sd.fields[f.name] = f
	}
	return sd.decode
}

type mapDecoder struct {
	kt, vt reflect.Type
}

func (md mapDecoder) decode(d *Decoder, v reflect.Value) error {
	if err := d.expectType(Map); err != nil {
		return err
	}
	ln, _, err := d.ReadUint()

	if err != nil {
		return err
	}

	t := v.Type()
	if v.IsNil() {
		v.Set(reflect.MakeMap(t))
	}
	kdec, vdec := typeDecoder(md.kt), typeDecoder(md.vt)
	for i, kt, vt := 0, t.Key(), t.Elem(); i < int(ln); i++ {
		key := reflect.New(kt).Elem()
		if err = kdec(d, key); err != nil {
			return err
		}
		if d.peekType() == Nil {
			v.SetMapIndex(key, reflect.Zero(vt))
			d.readType()
			continue
		}
		val := reflect.New(vt).Elem()
		if err = vdec(d, val); err != nil {
			return err
		}
		v.SetMapIndex(key, val)
	}
	return d.expectType(EOV)
}

func newMapDecoder(t reflect.Type) decoderFunc {
	md := mapDecoder{t.Key(), t.Elem()}
	return md.decode
}

type ptrDecoder struct {
	dec decoderFunc
}

func (pd ptrDecoder) decodeElem(d *Decoder, v reflect.Value) error {
	if v.IsNil() {
		v.Set(reflect.New(v.Type().Elem()))
	}
	return pd.dec(d, v.Elem())
}

func (pd ptrDecoder) decode(d *Decoder, v reflect.Value) error {
	if v.IsNil() {
		v.Set(reflect.New(v.Type().Elem()))
	}
	return pd.dec(d, v)
}

func newPtrDecoder(fn decoderFunc, elem bool) decoderFunc {
	pd := ptrDecoder{fn}
	if elem {
		return pd.decodeElem
	}
	return pd.decode
}

func addrDecoder(fn decoderFunc) decoderFunc {
	return func(d *Decoder, v reflect.Value) error {
		if v = v.Addr(); v.IsNil() {
			v.Set(reflect.New(v.Type().Elem().Elem()))
		}
		return d.Decode(v.Interface())
	}
}
<|endoftext|>