<commit_before>//   Copyright 2019 MSolution.IO
//
//   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 reports

import (
	"context"
	"database/sql"
	"fmt"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/360EntSecGroup-Skylar/excelize"
	"github.com/trackit/jsonlog"

	"github.com/trackit/trackit/aws"
	"github.com/trackit/trackit/aws/s3"
	"github.com/trackit/trackit/aws/usageReports/history"
	"github.com/trackit/trackit/costs/tags"
	"github.com/trackit/trackit/es"
	"github.com/trackit/trackit/users"
)

const tagsUsageReportSheetName = "Tags Report"

var tagsUsageReportModule = module{
	Name:          "Tags Usage Report",
	SheetName:     tagsUsageReportSheetName,
	ErrorName:     "tagsUsageReportError",
	GenerateSheet: generateTagsUsageReportSheet,
}

// generateTagsUsageReportSheet will generate a sheet with Tags usage report
// It will get data for given AWS account and for a given date
func generateTagsUsageReportSheet(ctx context.Context, aas []aws.AwsAccount, date time.Time, tx *sql.Tx, file *excelize.File) (err error) {
	if date.IsZero() {
		date, _ = history.GetHistoryDate()
	}
	return tagsUsageReportGenerateSheet(ctx, aas, date, tx, file)
}

func tagsUsageReportGenerateSheet(ctx context.Context, aas []aws.AwsAccount, date time.Time, tx *sql.Tx, file *excelize.File) (err error) {
	data, err := tagsUsageReportGetData(ctx, aas, date, tx)
	if err == nil {
		for key, report := range data {
			if err = tagsUsageReportInsertDataInSheet(ctx, file, key, report); err != nil {
				return
			}
		}
	}
	return
}

func getTagsKey(ctx context.Context, parsedParams tags.TagsValuesQueryParams) ([]string, error) {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	parsedParamsKeys := tags.TagsKeysQueryParams{
		AccountList: []string{},
		IndexList:   []string{},
		DateBegin:   parsedParams.DateBegin,
		DateEnd:     parsedParams.DateEnd,
	}
	_, keys, err := tags.GetTagsKeysWithParsedParams(ctx, parsedParamsKeys)
	if err != nil {
		logger.Error("Error when getting Tags Keys", err)
		return nil, err
	}
	return keys, nil
}

func tagsUsageReportGetData(ctx context.Context, aas []aws.AwsAccount, date time.Time, tx *sql.Tx) (reports tags.TagsValuesResponse, err error) {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	identities := getAwsIdentities(aas)
	user, err := users.GetUserWithId(tx, aas[0].UserId)
	if err != nil {
		return
	}
	parsedParams := tags.TagsValuesQueryParams{
		AccountList: []string{},
		IndexList:   []string{},
		DateBegin:   date,
		DateEnd:     time.Date(date.Year(), date.Month()+1, 0, 23, 59, 59, 999999999, date.Location()).UTC(),
		TagsKeys:    []string{},
		By:          "product",
		Detailed:    true,
	}
	parsedParams.AccountList = identities
	accountsAndIndexes, _, err := es.GetAccountsAndIndexes(parsedParams.AccountList, user, tx, s3.IndexPrefixLineItem)
	if err != nil {
		return nil, err
	}
	parsedParams.AccountList = accountsAndIndexes.Accounts
	parsedParams.IndexList = accountsAndIndexes.Indexes
	keys, err := getTagsKey(ctx, parsedParams)
	if err != nil {
		return nil, err
	}
	parsedParams.TagsKeys = keys
	logger.Debug("Getting Tags Usage Report for accounts", map[string]interface{}{
		"accounts": aas,
		"date":     date,
	})
	_, reports, err = tags.GetTagsValuesWithParsedParams(ctx, parsedParams)
	return
}

func tagsUsageReportInsertDataInSheet(ctx context.Context, file *excelize.File, key string, data []tags.TagsValues) (err error) {
	file.NewSheet(key)
	tagsUsageReportGenerateHeader(file, key)
	column := 2
	maxColumn := 2
	var totalCostCells []string
	var valueProductExist bool
	charColumn := 3
	for _, tag := range data {
		if tag.Tag == "" {
			tag.Tag = "no tag"
		}
		maxColumn, totalCostCells, valueProductExist = putProductDataInSheet(file, key, tag, maxColumn)
		if valueProductExist {
			cellsTagsCosts := cells{
				newCell(tag.Tag, "A"+strconv.Itoa(column)).mergeTo("A" + strconv.Itoa(maxColumn-1)),
				newFormula(fmt.Sprintf("SUM(%s)", strings.Join(totalCostCells, ",")), "F"+strconv.Itoa(column)).mergeTo("F" + strconv.Itoa(maxColumn-1)).addStyles("price"),
			}
			cellsTagsCosts.addStyles("borders", "centerText").setValues(file, key)
			cellsChart := cells{
				newCell(file.GetCellValue(key, fmt.Sprintf("A%s", strconv.Itoa(column))), "H"+strconv.Itoa(charColumn)),
				newFormula(fmt.Sprintf("=%s", "F"+strconv.Itoa(column)), "I"+strconv.Itoa(charColumn)).addStyles("price"),
			}
			cellsChart.addStyles("borders", "centerText").setValues(file, key)
			column = maxColumn
			charColumn++
		}
	}
	chartData := fmt.Sprintf(`{"name":"Costs","categories":"%s!$H$3:$H$%s","values":"%s!$I$3:$I$%s"}`, key, strconv.Itoa(charColumn - 1), key, strconv.Itoa(charColumn - 1))
	if err := generateLinearChart(ctx, file, key, chartData); err != nil {
		return err
	}
	return
}

func putProductDataInSheet(file *excelize.File, sheetName string, tag tags.TagsValues, maxColumn int) (int, []string, bool) {
	productColumn := maxColumn
	totalCostCells := make([]string, 0)
	valueProductExist := false
	for _, product := range tag.Items {
		productCostCells := make([]string, 0)
		valueTypeExist := false
		usageTypes := sortUsageTypesCosts(product.UsageTypes)
		for _, usageType := range usageTypes {
			if usageType.Cost != 0 {
				usageTypeCells := cells{
					newCell(usageType.UsageType, "C"+strconv.Itoa(maxColumn)),
					newCell(usageType.Cost, "D"+strconv.Itoa(maxColumn)).addStyles("price"),
				}
				usageTypeCells.addStyles("borders", "centerText").setValues(file, sheetName)
				maxColumn++
				productCostCells = append(productCostCells, "D"+strconv.Itoa(maxColumn))
				valueTypeExist = true
			}
		}
		if valueTypeExist {
			productCells := cells{
				newCell(product.Item, "B"+strconv.Itoa(productColumn)).mergeTo("B" + strconv.Itoa(maxColumn-1)),
				newFormula(fmt.Sprintf("SUM(%s)", strings.Join(productCostCells, ",")), "E"+strconv.Itoa(productColumn)).mergeTo("E" + strconv.Itoa(maxColumn-1)).addStyles("price"),
			}
			productCells.addStyles("borders", "centerText").setValues(file, sheetName)
			totalCostCells = append(totalCostCells, "E"+strconv.Itoa(productColumn))
			productColumn = maxColumn
			valueProductExist = true
		}
	}
	return maxColumn, totalCostCells, valueProductExist
}

func tagsUsageReportGenerateHeader(file *excelize.File, key string) {
	header := cells{
		newCell("Tags", "A1"),
		newCell("Products", "B1"),
		newCell("UsageTypes", "C1"),
		newCell("Costs", "D1"),
		newCell("Product Cost", "E1"),
		newCell("Total Cost", "F1"),
		newCell("Resume", "H1").mergeTo("I1"),
		newCell("Tags", "H2"),
		newCell("Total Cost", "I2"),
	}
	header.addStyles("borders", "bold", "centerText").setValues(file, key)
	columns := columnsWidth{
		newColumnWidth("A", 30),
		newColumnWidth("B", 20),
		newColumnWidth("C", 45),
		newColumnWidth("D", 15).toColumn("F"),
		newColumnWidth("H", 30),
	}
	columns.setValues(file, key)
	return
}

func generateLinearChart(ctx context.Context, file *excelize.File, sheetName, data string) error {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	err := file.AddChart(sheetName, "J1",fmt.Sprintf(
		`{"type":"bar",
				"series":[%s],
				"format":{
					"x_scale":2.0,"y_scale":2.0,
					"x_offset":15,
					"y_offset":10,
					"print_obj":true,
					"lock_aspect_ratio":false,
					"locked":false
				},
				"legend":{
					"position":"bottom",
					"show_legend_key":false
				},
				"title":{"name":"Cost per tag"},
				"plotarea":{
					"show_bubble_size":true,
					"show_cat_name":false,
					"show_leader_lines":false,
					"show_percent":true,
					"show_series_name":false,
					"show_val":false
				},
				"show_blanks_as":"gap"}`, data))
	if err != nil {
		logger.Error("Error when generating the chart", err)
	}
	return nil
}

func sortUsageTypesCosts(usageTypes []tags.ValueDetailed) (sortTypes []tags.ValueDetailed) {
	sortTypes = make([]tags.ValueDetailed, 0)
	toSort := make([]float64, 0)
	for _, usageType := range usageTypes {
		toSort = append(toSort, usageType.Cost)
	}
	sort.Float64s(toSort)
	sort.Sort(sort.Reverse(sort.Float64Slice(toSort)))
	for _, numberSort := range toSort {
		for _, usageType := range usageTypes {
			if alreadySort := checkTypeIsSort(usageType.UsageType, sortTypes); !alreadySort && usageType.Cost == numberSort {
				sortTypes = append(sortTypes, usageType)
				break
			}
		}
	}
	return
}

func checkTypeIsSort(usageType string, sortTypes []tags.ValueDetailed) bool {
	for _, sortType := range sortTypes {
		if sortType.UsageType == usageType {
			return true
		}
	}
	return false
}
<commit_msg>fix total product cost on tags spreadsheet<commit_after>//   Copyright 2019 MSolution.IO
//
//   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 reports

import (
	"context"
	"database/sql"
	"fmt"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/360EntSecGroup-Skylar/excelize"
	"github.com/trackit/jsonlog"

	"github.com/trackit/trackit/aws"
	"github.com/trackit/trackit/aws/s3"
	"github.com/trackit/trackit/aws/usageReports/history"
	"github.com/trackit/trackit/costs/tags"
	"github.com/trackit/trackit/es"
	"github.com/trackit/trackit/users"
)

const tagsUsageReportSheetName = "Tags Report"

var tagsUsageReportModule = module{
	Name:          "Tags Usage Report",
	SheetName:     tagsUsageReportSheetName,
	ErrorName:     "tagsUsageReportError",
	GenerateSheet: generateTagsUsageReportSheet,
}

// generateTagsUsageReportSheet will generate a sheet with Tags usage report
// It will get data for given AWS account and for a given date
func generateTagsUsageReportSheet(ctx context.Context, aas []aws.AwsAccount, date time.Time, tx *sql.Tx, file *excelize.File) (err error) {
	if date.IsZero() {
		date, _ = history.GetHistoryDate()
	}
	return tagsUsageReportGenerateSheet(ctx, aas, date, tx, file)
}

func tagsUsageReportGenerateSheet(ctx context.Context, aas []aws.AwsAccount, date time.Time, tx *sql.Tx, file *excelize.File) (err error) {
	data, err := tagsUsageReportGetData(ctx, aas, date, tx)
	if err == nil {
		for key, report := range data {
			if err = tagsUsageReportInsertDataInSheet(ctx, file, key, report); err != nil {
				return
			}
		}
	}
	return
}

func getTagsKey(ctx context.Context, parsedParams tags.TagsValuesQueryParams) ([]string, error) {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	parsedParamsKeys := tags.TagsKeysQueryParams{
		AccountList: []string{},
		IndexList:   []string{},
		DateBegin:   parsedParams.DateBegin,
		DateEnd:     parsedParams.DateEnd,
	}
	_, keys, err := tags.GetTagsKeysWithParsedParams(ctx, parsedParamsKeys)
	if err != nil {
		logger.Error("Error when getting Tags Keys", err)
		return nil, err
	}
	return keys, nil
}

func tagsUsageReportGetData(ctx context.Context, aas []aws.AwsAccount, date time.Time, tx *sql.Tx) (reports tags.TagsValuesResponse, err error) {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	identities := getAwsIdentities(aas)
	user, err := users.GetUserWithId(tx, aas[0].UserId)
	if err != nil {
		return
	}
	parsedParams := tags.TagsValuesQueryParams{
		AccountList: []string{},
		IndexList:   []string{},
		DateBegin:   date,
		DateEnd:     time.Date(date.Year(), date.Month()+1, 0, 23, 59, 59, 999999999, date.Location()).UTC(),
		TagsKeys:    []string{},
		By:          "product",
		Detailed:    true,
	}
	parsedParams.AccountList = identities
	accountsAndIndexes, _, err := es.GetAccountsAndIndexes(parsedParams.AccountList, user, tx, s3.IndexPrefixLineItem)
	if err != nil {
		return nil, err
	}
	parsedParams.AccountList = accountsAndIndexes.Accounts
	parsedParams.IndexList = accountsAndIndexes.Indexes
	keys, err := getTagsKey(ctx, parsedParams)
	if err != nil {
		return nil, err
	}
	parsedParams.TagsKeys = keys
	logger.Debug("Getting Tags Usage Report for accounts", map[string]interface{}{
		"accounts": aas,
		"date":     date,
	})
	_, reports, err = tags.GetTagsValuesWithParsedParams(ctx, parsedParams)
	return
}

func tagsUsageReportInsertDataInSheet(ctx context.Context, file *excelize.File, key string, data []tags.TagsValues) (err error) {
	file.NewSheet(key)
	tagsUsageReportGenerateHeader(file, key)
	column := 2
	maxColumn := 2
	var totalCostCells []string
	var valueProductExist bool
	charColumn := 3
	for _, tag := range data {
		if tag.Tag == "" {
			tag.Tag = "no tag"
		}
		maxColumn, totalCostCells, valueProductExist = putProductDataInSheet(file, key, tag, maxColumn)
		if valueProductExist {
			cellsTagsCosts := cells{
				newCell(tag.Tag, "A"+strconv.Itoa(column)).mergeTo("A" + strconv.Itoa(maxColumn-1)),
				newFormula(fmt.Sprintf("SUM(%s)", strings.Join(totalCostCells, ",")), "F"+strconv.Itoa(column)).mergeTo("F" + strconv.Itoa(maxColumn-1)).addStyles("price"),
			}
			cellsTagsCosts.addStyles("borders", "centerText").setValues(file, key)
			cellsChart := cells{
				newCell(file.GetCellValue(key, fmt.Sprintf("A%s", strconv.Itoa(column))), "H"+strconv.Itoa(charColumn)),
				newFormula(fmt.Sprintf("=%s", "F"+strconv.Itoa(column)), "I"+strconv.Itoa(charColumn)).addStyles("price"),
			}
			cellsChart.addStyles("borders", "centerText").setValues(file, key)
			column = maxColumn
			charColumn++
		}
	}
	chartData := fmt.Sprintf(`{"name":"Costs","categories":"%s!$H$3:$H$%s","values":"%s!$I$3:$I$%s"}`, key, strconv.Itoa(charColumn - 1), key, strconv.Itoa(charColumn - 1))
	if err := generateLinearChart(ctx, file, key, chartData); err != nil {
		return err
	}
	return
}

func putProductDataInSheet(file *excelize.File, sheetName string, tag tags.TagsValues, maxColumn int) (int, []string, bool) {
	productColumn := maxColumn
	totalCostCells := make([]string, 0)
	valueProductExist := false
	for _, product := range tag.Items {
		productCostCells := make([]string, 0)
		valueTypeExist := false
		usageTypes := sortUsageTypesCosts(product.UsageTypes)
		for _, usageType := range usageTypes {
			if usageType.Cost != 0 {
				usageTypeCells := cells{
					newCell(usageType.UsageType, "C"+strconv.Itoa(maxColumn)),
					newCell(usageType.Cost, "D"+strconv.Itoa(maxColumn)).addStyles("price"),
				}
				usageTypeCells.addStyles("borders", "centerText").setValues(file, sheetName)
				maxColumn++
				productCostCells = append(productCostCells, "D"+strconv.Itoa(maxColumn-1))
				valueTypeExist = true
			}
		}
		if valueTypeExist {
			productCells := cells{
				newCell(product.Item, "B"+strconv.Itoa(productColumn)).mergeTo("B" + strconv.Itoa(maxColumn-1)),
				newFormula(fmt.Sprintf("SUM(%s)", strings.Join(productCostCells, ",")), "E"+strconv.Itoa(productColumn)).mergeTo("E" + strconv.Itoa(maxColumn-1)).addStyles("price"),
			}
			fmt.Printf("productcolumn = %v, maxColumn = %v\n", productColumn, maxColumn)
			productCells.addStyles("borders", "centerText").setValues(file, sheetName)
			totalCostCells = append(totalCostCells, "E"+strconv.Itoa(productColumn))
			productColumn = maxColumn
			valueProductExist = true
		}
	}
	return maxColumn, totalCostCells, valueProductExist
}

func tagsUsageReportGenerateHeader(file *excelize.File, key string) {
	header := cells{
		newCell("Tags", "A1"),
		newCell("Products", "B1"),
		newCell("UsageTypes", "C1"),
		newCell("Costs", "D1"),
		newCell("Product Cost", "E1"),
		newCell("Total Cost", "F1"),
		newCell("Resume", "H1").mergeTo("I1"),
		newCell("Tags", "H2"),
		newCell("Total Cost", "I2"),
	}
	header.addStyles("borders", "bold", "centerText").setValues(file, key)
	columns := columnsWidth{
		newColumnWidth("A", 30),
		newColumnWidth("B", 20),
		newColumnWidth("C", 45),
		newColumnWidth("D", 15).toColumn("F"),
		newColumnWidth("H", 30),
	}
	columns.setValues(file, key)
	return
}

func generateLinearChart(ctx context.Context, file *excelize.File, sheetName, data string) error {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	err := file.AddChart(sheetName, "J1",fmt.Sprintf(
		`{"type":"bar",
				"series":[%s],
				"format":{
					"x_scale":2.0,"y_scale":2.0,
					"x_offset":15,
					"y_offset":10,
					"print_obj":true,
					"lock_aspect_ratio":false,
					"locked":false
				},
				"legend":{
					"position":"bottom",
					"show_legend_key":false
				},
				"title":{"name":"Cost per tag"},
				"plotarea":{
					"show_bubble_size":true,
					"show_cat_name":false,
					"show_leader_lines":false,
					"show_percent":true,
					"show_series_name":false,
					"show_val":false
				},
				"show_blanks_as":"gap"}`, data))
	if err != nil {
		logger.Error("Error when generating the chart", err)
	}
	return nil
}

func sortUsageTypesCosts(usageTypes []tags.ValueDetailed) (sortTypes []tags.ValueDetailed) {
	sortTypes = make([]tags.ValueDetailed, 0)
	toSort := make([]float64, 0)
	for _, usageType := range usageTypes {
		toSort = append(toSort, usageType.Cost)
	}
	sort.Float64s(toSort)
	sort.Sort(sort.Reverse(sort.Float64Slice(toSort)))
	for _, numberSort := range toSort {
		for _, usageType := range usageTypes {
			if alreadySort := checkTypeIsSort(usageType.UsageType, sortTypes); !alreadySort && usageType.Cost == numberSort {
				sortTypes = append(sortTypes, usageType)
				break
			}
		}
	}
	return
}

func checkTypeIsSort(usageType string, sortTypes []tags.ValueDetailed) bool {
	for _, sortType := range sortTypes {
		if sortType.UsageType == usageType {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/amahi/spdy"
	"golang.org/x/net/http2"
	"io/ioutil"
	"jgcrypto/tls"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"
)

// in returns true is s is one of the strings in list
func in(list []string, s string) bool {
	for _, t := range list {
		if s == t {
			return true
		}
	}

	return false
}

// tri captures a tri-state. The value of yesno is true only is ran is
// true
type tri struct {
	ran   bool
	yesno bool
}

func (t tri) String() string {
	switch {
	case !t.ran:
		return "-"
	case t.yesno:
		return "t"
	case !t.yesno:
		return "f"
	}

	// Should not be reached ever

	return "!"
}

// site is a web site identified by its DNS name along with the state
// of various tests performed on the site. The states are only filled
// in if test() is called on the site.
type site struct {
	name string        // DNS name of the web site

	resolves       tri // Whether the name resolves
	port443Open    tri // Whether port 443 is open via TCP
	tlsWorks       tri // Whether a TLS connection works
	httpsWorks     tri // Whether an HTTPS request works using HTTP/1.1
	spdyAnnounced  tri // Whether spdy/3.1 is announced using NPN
	http2Announced tri // Whether h2 is announced using NPN
	spdyWorks      tri // Whether a SPDY/3.1 request works
	http2Works     tri // Whether an HTTP/2 request works

	npn            []string // List of protocols advertised by server using NPN
}

// test tests a site to see if it supports various protocols
func (s *site) test(l *os.File) {

	// Check the name resolves, give up if it does not

	s.resolves.ran = true
	_, err := net.LookupHost(s.name)
	if err != nil {
		s.logf(l, "Error resolving name: %s", err)
		s.resolves.yesno = false
		return
	}
	s.resolves.yesno = true

	// See if port 443 is open, give up if it is not

	hostPort := net.JoinHostPort(s.name, "443")

	s.port443Open.ran = true
	c, err := net.DialTimeout("tcp", hostPort, 2 * time.Second)
	if err != nil {
		s.logf(l, "TCP dial to port 443 failed: %s", err)
		s.port443Open.yesno = false
		return
	}
	c.Close()
	s.port443Open.yesno = true

	// See if TLS works and if SPDY/3.1 or HTTP/2 offered. Give up if
	// TLS does not work. If the name doesn't work try adding www.  to
	// se if it's a TLS certificate error.

	s.tlsWorks.ran = true
	config := &tls.Config{}
	config.ServerName = s.name
	tc, err := tls.Dial("tcp", hostPort, config)
	if err != nil {
		s.name = "www." + s.name
		config.ServerName = s.name
		hostPort = net.JoinHostPort(s.name, "443")
		tc, err = tls.Dial("tcp", hostPort, config)
		if err != nil {
			s.logf(l, "Error performing TLS connection: %s", err)
			s.tlsWorks.yesno = false
			return
		}
	}
	s.tlsWorks.yesno = true

	// Retrieve the list of NPN offered protocols and check for
	// spdy/3.1 and h2

	cs := tc.ConnectionState()
	s.npn = cs.OfferedProtocols
	s.spdyAnnounced.ran = true
	s.spdyAnnounced.yesno = in(cs.OfferedProtocols, "spdy/3.1")
	s.http2Announced.ran = true
	s.http2Announced.yesno = in(cs.OfferedProtocols, "h2")
	tc.Close()

	// See if HTTPS works by performing GET /

	s.httpsWorks.ran = true
	resp, err := http.Get("https://" + s.name)
	if err != nil {
		s.logf(l, "HTTP request failed: %s", err)
		return
	}
	if resp != nil && resp.Body != nil {
		_, _ = ioutil.ReadAll(resp.Body)
		resp.Body.Close()
	}
	s.httpsWorks.yesno = err == nil

	// See if SPDY works by doing GET / over SPDY

	if s.spdyAnnounced.yesno {
		s.spdyWorks.ran = true
		conf := &tls.Config{}
		conf.NextProtos = []string{"spdy/3.1"}
		spdyC, err := tls.Dial("tcp", hostPort, conf)
		if err == nil {
			defer spdyC.Close()
			cs = spdyC.ConnectionState()
			if cs.NegotiatedProtocol != "spdy/3.1" {
				s.logf(l, "NegotiatedProtocol not spdy/3.1: %s",
					cs.NegotiatedProtocol)
			} else {
				sc, err := spdy.NewClientConn(spdyC)
				if err == nil {
					defer sc.Close()
					req, _ := http.NewRequest("GET", "https://"+s.name, nil)
					resp, err := sc.Do(req)
					if err != nil {
						s.logf(l, "Failed to do SPDY request: %s", err)
					}
					s.spdyWorks.yesno = err == nil
					if resp != nil && resp.Body != nil {
						_, _ = ioutil.ReadAll(resp.Body)
						resp.Body.Close()
					}
				} else {
					s.logf(l, "Failed create SPDY client connection: %s", err)
				}
			}
		} else {
			s.logf(l, "Failed to dial port 443 for SPDY: %s", err)
		}
	}

	// See if HTTP/2 works by doing GET / using HTTP/2

	if s.http2Announced.yesno {
		s.http2Works.ran = true
		conf := &tls.Config{}
		conf.NextProtos = []string{"h2"}
		h2C, err := tls.Dial("tcp", hostPort, conf)
		if err == nil {
			defer h2C.Close()
			cs = h2C.ConnectionState()
			if cs.NegotiatedProtocol != "h2" {
				s.logf(l, "NegotiatedProtocol not h2: %s",
					cs.NegotiatedProtocol)
			} else {
				h2t := &http2.Transport{}
				h2c, err := h2t.NewClientConn(h2C)
				if err == nil {
					req, _ := http.NewRequest("GET", "https://"+s.name, nil)
					resp, err := h2c.RoundTrip(req)
					if err != nil {
						s.logf(l, "HTTP/2 RoundTrip failed: %s", err)
					}
					s.http2Works.yesno = err == nil
					if resp != nil && resp.Body != nil {
						_, _ = ioutil.ReadAll(resp.Body)
						resp.Body.Close()
					}
				} else {
					s.logf(l, "Failed create HTTP/2 client connection: %s", err)
				}
				h2t.CloseIdleConnections()
			}
		} else {
			s.logf(l, "Failed to dial port 443 for HTTP/2: %s", err)
		}
	}
}

// logf writes to the log file prefixing with the name of the site
// being logged
func (s *site) logf(f *os.File, format string, a ...interface{}) {
	if f != nil {
		fmt.Fprintf(f, fmt.Sprintf(s.name + ": " + format + "\n", a...))
	}
}

// fields returns the list of fields that String() will return for a
// site
func (s *site) fields() string {
	return "name,resolves,port443Open,tlsWorks,httpsWorks,spdyAnnounced,http2Announced,spdyWorks,http2Works,npn"
}

func (s *site) String() string {
	return fmt.Sprintf("%s,%s,%s,%s,%s,%s,%s,%s,%s,%s",
		s.name, s.resolves, s.port443Open, s.tlsWorks, s.httpsWorks,
		s.spdyAnnounced, s.http2Announced, s.spdyWorks, s.http2Works,
		strings.Join(s.npn, " "))
}

var wg sync.WaitGroup

func worker(work, result chan *site, l *os.File) {
	for s := range work {
		s.test(l)
		result <- s
	}
	wg.Done()
}

func writer(result chan *site, stop chan struct{}, fields bool) {
	first := true
	for s := range result {
		if fields && first {
			fmt.Printf("%s\n", s.fields())
			first = false
		}

		fmt.Printf("%s\n", s)
	}
	close(stop)
}

func main() {

	// The SPDY library is chatty so discard its log statements

	spdy.SetLog(ioutil.Discard)

	fields := flag.Bool("fields", false,
		"If set outputs a header line containing field names")
	workers := flag.Int("workers", 10, "Number of concurrent workers")
	log := flag.String("log", "", "File to write log information to")
	flag.Parse()

	if *workers < 1 {
		fmt.Printf("-workers must be a positive number\n")
		return
	}

	var l *os.File
	var err error
	if *log != "" {
		if l, err = os.Create(*log); err != nil {
			fmt.Printf("Failed to create log file %s: %s\n", *log, err)
			return
		}
		defer l.Close()
	}

	work := make(chan *site)
	result := make(chan *site)
	stop := make(chan struct{})

	go writer(result, stop, *fields)

	for i := 0; i < *workers; i++ {
		wg.Add(1)
		go worker(work, result, l)
	}

	scan := bufio.NewScanner(os.Stdin)
	for scan.Scan() {
		work <- &site{name: scan.Text()}
	}

	close(work)
	wg.Wait()
	close(result)
	<-stop

	if scan.Err() != nil {
		fmt.Printf("Error reading input: %s\n", scan.Err())
		return
	}
}
<commit_msg>Add timeouts to dials, deadline to I/O and use RoundTrip instead of Get<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/amahi/spdy"
	"golang.org/x/net/http2"
	"io/ioutil"
	"jgcrypto/tls"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"
)

// in returns true is s is one of the strings in list
func in(list []string, s string) bool {
	for _, t := range list {
		if s == t {
			return true
		}
	}

	return false
}

// tri captures a tri-state. The value of yesno is true only is ran is
// true
type tri struct {
	ran   bool
	yesno bool
}

func (t tri) String() string {
	switch {
	case !t.ran:
		return "-"
	case t.yesno:
		return "t"
	case !t.yesno:
		return "f"
	}

	// Should not be reached ever

	return "!"
}

// site is a web site identified by its DNS name along with the state
// of various tests performed on the site. The states are only filled
// in if test() is called on the site.
type site struct {
	name string        // DNS name of the web site

	resolves       tri // Whether the name resolves
	port443Open    tri // Whether port 443 is open via TCP
	tlsWorks       tri // Whether a TLS connection works
	httpsWorks     tri // Whether an HTTPS request works using HTTP/1.1
	spdyAnnounced  tri // Whether spdy/3.1 is announced using NPN
	http2Announced tri // Whether h2 is announced using NPN
	spdyWorks      tri // Whether a SPDY/3.1 request works
	http2Works     tri // Whether an HTTP/2 request works

	npn            []string // List of protocols advertised by server using NPN
}

// test tests a site to see if it supports various protocols
func (s *site) test(l *os.File) {

	// Check the name resolves, give up if it does not

	s.resolves.ran = true
	_, err := net.LookupHost(s.name)
	if err != nil {
		s.logf(l, "Error resolving name: %s", err)
		s.resolves.yesno = false
		return
	}
	s.resolves.yesno = true

	// See if port 443 is open, give up if it is not

	hostPort := net.JoinHostPort(s.name, "443")

	s.port443Open.ran = true
	c, err := net.DialTimeout("tcp", hostPort, 2 * time.Second)
	if err != nil {
		s.logf(l, "TCP dial to port 443 failed: %s", err)
		s.port443Open.yesno = false
		return
	}
	c.Close()
	s.port443Open.yesno = true

	// See if TLS works and if SPDY/3.1 or HTTP/2 offered. Give up if
	// TLS does not work. If the name doesn't work try adding www.  to
	// se if it's a TLS certificate error.

	s.tlsWorks.ran = true
	config := &tls.Config{}
	config.ServerName = s.name
	timeoutDialer := &net.Dialer{}
	timeoutDialer.Timeout = 2 * time.Second
	tc, err := tls.DialWithDialer(timeoutDialer, "tcp", hostPort, config)
	if err != nil {
		s.name = "www." + s.name
		config.ServerName = s.name
		hostPort = net.JoinHostPort(s.name, "443")
		tc, err = tls.DialWithDialer(timeoutDialer, "tcp", hostPort, config)
		if err != nil {
			s.logf(l, "Error performing TLS connection: %s", err)
			s.tlsWorks.yesno = false
			return
		}
	}
	s.tlsWorks.yesno = true

	// Retrieve the list of NPN offered protocols and check for
	// spdy/3.1 and h2

	cs := tc.ConnectionState()
	s.npn = cs.OfferedProtocols
	s.spdyAnnounced.ran = true
	s.spdyAnnounced.yesno = in(cs.OfferedProtocols, "spdy/3.1")
	s.http2Announced.ran = true
	s.http2Announced.yesno = in(cs.OfferedProtocols, "h2")
	tc.Close()

	// See if HTTPS works by performing GET /

	s.httpsWorks.ran = true
	httpTransport := &http.Transport{}
	req, _ := http.NewRequest("GET", "https://" + s.name, nil)
	resp, err := httpTransport.RoundTrip(req)
	if err != nil {
		s.logf(l, "HTTP request failed: %s", err)
		return
	}
	if resp != nil && resp.Body != nil {
		_, _ = ioutil.ReadAll(resp.Body)
		resp.Body.Close()
	}
	httpTransport.CloseIdleConnections()
	s.httpsWorks.yesno = err == nil

	// See if SPDY works by doing GET / over SPDY

	if s.spdyAnnounced.yesno {
		s.spdyWorks.ran = true
		conf := &tls.Config{}
		conf.NextProtos = []string{"spdy/3.1"}
		spdyC, err := tls.DialWithDialer(timeoutDialer, "tcp", hostPort, conf)
		if err == nil {
			spdyC.SetDeadline(time.Now().Add(time.Minute))
			defer spdyC.Close()
			cs = spdyC.ConnectionState()
			if cs.NegotiatedProtocol != "spdy/3.1" {
				s.logf(l, "NegotiatedProtocol not spdy/3.1: %s",
					cs.NegotiatedProtocol)
			} else {
				sc, err := spdy.NewClientConn(spdyC)
				if err == nil {
					defer sc.Close()
					req, _ := http.NewRequest("GET", "https://"+s.name, nil)
					resp, err := sc.Do(req)
					if err != nil {
						s.logf(l, "Failed to do SPDY request: %s", err)
					}
					s.spdyWorks.yesno = err == nil
					if resp != nil && resp.Body != nil {
						_, _ = ioutil.ReadAll(resp.Body)
						resp.Body.Close()
					}
				} else {
					s.logf(l, "Failed create SPDY client connection: %s", err)
				}
			}
		} else {
			s.logf(l, "Failed to dial port 443 for SPDY: %s", err)
		}
	}

	// See if HTTP/2 works by doing GET / using HTTP/2

	if s.http2Announced.yesno {
		s.http2Works.ran = true
		conf := &tls.Config{}
		conf.NextProtos = []string{"h2"}
		h2C, err := tls.DialWithDialer(timeoutDialer, "tcp", hostPort, conf)
		if err == nil {
			h2C.SetDeadline(time.Now().Add(time.Minute))
			defer h2C.Close()
			cs = h2C.ConnectionState()
			if cs.NegotiatedProtocol != "h2" {
				s.logf(l, "NegotiatedProtocol not h2: %s",
					cs.NegotiatedProtocol)
			} else {
				h2t := &http2.Transport{}
				h2c, err := h2t.NewClientConn(h2C)
				if err == nil {
					req, _ := http.NewRequest("GET", "https://"+s.name, nil)
					resp, err := h2c.RoundTrip(req)
					if err != nil {
						s.logf(l, "HTTP/2 RoundTrip failed: %s", err)
					}
					s.http2Works.yesno = err == nil
					if resp != nil && resp.Body != nil {
						_, _ = ioutil.ReadAll(resp.Body)
						resp.Body.Close()
					}
				} else {
					s.logf(l, "Failed create HTTP/2 client connection: %s", err)
				}
				h2t.CloseIdleConnections()
			}
		} else {
			s.logf(l, "Failed to dial port 443 for HTTP/2: %s", err)
		}
	}
}

// logf writes to the log file prefixing with the name of the site
// being logged
func (s *site) logf(f *os.File, format string, a ...interface{}) {
	if f != nil {
		fmt.Fprintf(f, fmt.Sprintf(s.name + ": " + format + "\n", a...))
	}
}

// fields returns the list of fields that String() will return for a
// site
func (s *site) fields() string {
	return "name,resolves,port443Open,tlsWorks,httpsWorks,spdyAnnounced,http2Announced,spdyWorks,http2Works,npn"
}

func (s *site) String() string {
	return fmt.Sprintf("%s,%s,%s,%s,%s,%s,%s,%s,%s,%s",
		s.name, s.resolves, s.port443Open, s.tlsWorks, s.httpsWorks,
		s.spdyAnnounced, s.http2Announced, s.spdyWorks, s.http2Works,
		strings.Join(s.npn, " "))
}

var wg sync.WaitGroup

func worker(work, result chan *site, l *os.File) {
	for s := range work {
		s.test(l)
		result <- s
	}
	wg.Done()
}

func writer(result chan *site, stop chan struct{}, fields bool) {
	first := true
	for s := range result {
		if fields && first {
			fmt.Printf("%s\n", s.fields())
			first = false
		}

		fmt.Printf("%s\n", s)
	}
	close(stop)
}

func main() {

	// The SPDY library is chatty so discard its log statements

	spdy.SetLog(ioutil.Discard)

	fields := flag.Bool("fields", false,
		"If set outputs a header line containing field names")
	workers := flag.Int("workers", 10, "Number of concurrent workers")
	log := flag.String("log", "", "File to write log information to")
	flag.Parse()

	if *workers < 1 {
		fmt.Printf("-workers must be a positive number\n")
		return
	}

	var l *os.File
	var err error
	if *log != "" {
		if l, err = os.Create(*log); err != nil {
			fmt.Printf("Failed to create log file %s: %s\n", *log, err)
			return
		}
		defer l.Close()
	}

	work := make(chan *site)
	result := make(chan *site)
	stop := make(chan struct{})

	go writer(result, stop, *fields)

	for i := 0; i < *workers; i++ {
		wg.Add(1)
		go worker(work, result, l)
	}

	scan := bufio.NewScanner(os.Stdin)
	for scan.Scan() {
		work <- &site{name: scan.Text()}
	}

	close(work)
	wg.Wait()
	close(result)
	<-stop

	if scan.Err() != nil {
		fmt.Printf("Error reading input: %s\n", scan.Err())
		return
	}
}
<|endoftext|>


<commit_before>/*-
Copyright (c) 2016, Jörg Pernfuß <code.jpe@gmail.com>
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

package auth

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"net"
	"net/http"
	"time"

	"github.com/mjolnir42/scrypth64"
)

// Token is the data passed between client and server to
// authenticate the client and issue a token for it that can be used
// as HTTP Basic Auth password.
type Token struct {
	UserName  string `json:"username"`
	Password  string `json:"password,omitempty"`
	Token     string `json:"token,omitempty"`
	ValidFrom string `json:"validFrom,omitempty"`
	ExpiresAt string `json:"expiresAt,omitempty"`
	Salt      string `json:"-"`
	SourceIP  net.IP `json:"-"`
}

// NewToken returns an empty Token
func NewToken() *Token {
	return &Token{}
}

// Generate verifies a the embedded credentials in Token and
// issues a new token to be returned to the user. Calling
// GenerateToken consumes the embedded password regardless of outcome.
// Returns ErrAuth if the password is incorrect.
func (t *Token) Generate(mcf scrypth64.Mcf, key, seed []byte) error {
	var (
		err                 error
		ok                  bool
		stuntSeed, stuntKey []byte
	)
	// do not generate Tokens with an incorrectly configured seed or
	// key
	if seed == nil || key == nil || len(seed) == 0 || len(key) == 0 {
		err = ErrInput
		goto fail
	}

	// net.IP is a bit of a pain since an unset net.IP does not equal
	// nil, and .String() prints `<nil>`
	if t.SourceIP.Equal(net.IP{}) {
		err = ErrInput
		goto fail
	}

	// Username and password have to be set
	if t.UserName == "" || t.Password == "" {
		err = ErrInput
		goto fail
	}

	// stunt data
	stuntSeed = bytes.Repeat([]byte{0x0F}, len(seed))
	stuntKey = bytes.Repeat([]byte{0xAB}, len(key))

	// start of timing critical path
	if ok, err = t.comparePassword(mcf); !ok || err != nil {
		// spend some quality time computing garbage
		t.mixToken(stuntKey, stuntSeed)
		goto fail
	}

	if err = t.mixToken(key, seed); err != nil {
		goto fail
	}

	t.zeroPassword()
	return nil

fail:
	t.zeroPassword()
	if err != nil {
		return err
	}
	return ErrAuth
}

// zeroPassword ensures the Password field is set to the zero value
func (t *Token) zeroPassword() {
	t.Password = ""
}

// comparePassword is used to verify the user supplied password
// against a scrypth64.Mcf
func (t *Token) comparePassword(mcf scrypth64.Mcf) (bool, error) {
	return scrypth64.Verify(t.Password, mcf)
}

// mixToken generates a new password token
func (t *Token) mixToken(key, seed []byte) error {
	var (
		salt, bin, btime []byte
		err              error
		valid, expires   time.Time
	)
	// expiry time
	valid = time.Now().UTC()
	expires = valid.Add(
		time.Duration(TokenExpirySeconds) * time.Second,
	).UTC()
	if btime, err = expires.MarshalBinary(); err != nil {
		goto fail
	}
	t.ValidFrom = valid.Format(rfc3339Milli)
	t.ExpiresAt = expires.Format(rfc3339Milli)

	// whiteout unstable subsecond timestamp part with "random" value
	copy(btime[9:], []byte{0xde, 0xad, 0xca, 0xfe})

	// add random salt
	salt = make([]byte, 16)
	if _, err = rand.Read(salt); err != nil {
		goto fail
	}
	t.Salt = hex.EncodeToString(salt)

	// compute the token
	bin = computeToken(
		[]byte(t.UserName),
		key,
		seed,
		btime,
		salt,
		[]byte(t.SourceIP.String()),
	)
	t.Token = hex.EncodeToString(bin)

	return nil

fail:
	return err
}

// SetIPAddress records the client's IP address
func (t *Token) SetIPAddress(r *http.Request) {
	t.SetIPAddressString(r.RemoteAddr)
}

// SetIPAddressString records the client's IP address
func (t *Token) SetIPAddressString(addr string) {
	t.SourceIP = net.ParseIP(extractAddress(addr))
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>Expose token.SetIPAddressExtractedString<commit_after>/*-
Copyright (c) 2016, Jörg Pernfuß <code.jpe@gmail.com>
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

package auth

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"net"
	"net/http"
	"time"

	"github.com/mjolnir42/scrypth64"
)

// Token is the data passed between client and server to
// authenticate the client and issue a token for it that can be used
// as HTTP Basic Auth password.
type Token struct {
	UserName  string `json:"username"`
	Password  string `json:"password,omitempty"`
	Token     string `json:"token,omitempty"`
	ValidFrom string `json:"validFrom,omitempty"`
	ExpiresAt string `json:"expiresAt,omitempty"`
	Salt      string `json:"-"`
	SourceIP  net.IP `json:"-"`
}

// NewToken returns an empty Token
func NewToken() *Token {
	return &Token{}
}

// Generate verifies a the embedded credentials in Token and
// issues a new token to be returned to the user. Calling
// GenerateToken consumes the embedded password regardless of outcome.
// Returns ErrAuth if the password is incorrect.
func (t *Token) Generate(mcf scrypth64.Mcf, key, seed []byte) error {
	var (
		err                 error
		ok                  bool
		stuntSeed, stuntKey []byte
	)
	// do not generate Tokens with an incorrectly configured seed or
	// key
	if seed == nil || key == nil || len(seed) == 0 || len(key) == 0 {
		err = ErrInput
		goto fail
	}

	// net.IP is a bit of a pain since an unset net.IP does not equal
	// nil, and .String() prints `<nil>`
	if t.SourceIP.Equal(net.IP{}) {
		err = ErrInput
		goto fail
	}

	// Username and password have to be set
	if t.UserName == "" || t.Password == "" {
		err = ErrInput
		goto fail
	}

	// stunt data
	stuntSeed = bytes.Repeat([]byte{0x0F}, len(seed))
	stuntKey = bytes.Repeat([]byte{0xAB}, len(key))

	// start of timing critical path
	if ok, err = t.comparePassword(mcf); !ok || err != nil {
		// spend some quality time computing garbage
		t.mixToken(stuntKey, stuntSeed)
		goto fail
	}

	if err = t.mixToken(key, seed); err != nil {
		goto fail
	}

	t.zeroPassword()
	return nil

fail:
	t.zeroPassword()
	if err != nil {
		return err
	}
	return ErrAuth
}

// zeroPassword ensures the Password field is set to the zero value
func (t *Token) zeroPassword() {
	t.Password = ""
}

// comparePassword is used to verify the user supplied password
// against a scrypth64.Mcf
func (t *Token) comparePassword(mcf scrypth64.Mcf) (bool, error) {
	return scrypth64.Verify(t.Password, mcf)
}

// mixToken generates a new password token
func (t *Token) mixToken(key, seed []byte) error {
	var (
		salt, bin, btime []byte
		err              error
		valid, expires   time.Time
	)
	// expiry time
	valid = time.Now().UTC()
	expires = valid.Add(
		time.Duration(TokenExpirySeconds) * time.Second,
	).UTC()
	if btime, err = expires.MarshalBinary(); err != nil {
		goto fail
	}
	t.ValidFrom = valid.Format(rfc3339Milli)
	t.ExpiresAt = expires.Format(rfc3339Milli)

	// whiteout unstable subsecond timestamp part with "random" value
	copy(btime[9:], []byte{0xde, 0xad, 0xca, 0xfe})

	// add random salt
	salt = make([]byte, 16)
	if _, err = rand.Read(salt); err != nil {
		goto fail
	}
	t.Salt = hex.EncodeToString(salt)

	// compute the token
	bin = computeToken(
		[]byte(t.UserName),
		key,
		seed,
		btime,
		salt,
		[]byte(t.SourceIP.String()),
	)
	t.Token = hex.EncodeToString(bin)

	return nil

fail:
	return err
}

// SetIPAddress records the client's IP address
func (t *Token) SetIPAddress(r *http.Request) {
	t.SetIPAddressString(r.RemoteAddr)
}

// SetIPAddressString records the client's IP address
func (t *Token) SetIPAddressString(addr string) {
	t.SetIPAddressExtractedString(extractAddress(addr))
}

// SetIPAddressExtractedString records the client's IP address
func (t *Token) SetIPAddressExtractedString(addr string) {
	t.SourceIP = net.ParseIP(addr)
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<|endoftext|>


<commit_before>package cli

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	uuid "github.com/nu7hatch/gouuid"
	"io/ioutil"
	"net/http"
	"os/user"
	"strings"
	"time"
)

// var authSite = "https://koding.com"

var authSite = "http://localhost:3020"

type Kd struct {
	ID         string
	PublicKey  string
	PrivateKey string
	Username   string
}

func NewKd() *Kd {
	id, _ := uuid.NewV4()
	return &Kd{
		ID: "machineID-" + id.String(),
	}
}

func (r *Kd) Definition() string {
	return "Registers the kite system to kontrol"
}

func (r *Kd) Exec() error {
	k := NewKd()
	err := k.Register()
	if err != nil {
		return errors.New(fmt.Sprint("\nregister error: %s\n", err.Error()))
	}
	return nil
}

func (k *Kd) Register() error {
	err := k.GetConfig()
	if err != nil {
		return err
	}

	registerUrl := fmt.Sprintf("%s/-/auth/register/%s/%s", authSite, k.ID, k.PublicKey)
	checkUrl := fmt.Sprintf("%s/-/auth/check/%s", authSite, k.PublicKey)

	fmt.Printf("Please open the following url for authentication:\n\n")
	// PrintBox(registerUrl)
	fmt.Println(registerUrl)
	fmt.Printf("\nwaiting . ")

	// check the result every two seconds
	ticker := time.NewTicker(time.Second * 2).C

	// wait for three minutes, if not successfull abort it
	timeout := time.After(time.Minute * 3)

	for {
		select {
		case <-ticker:
			resp, err := http.Get(checkUrl)
			if err != nil {
				return err
			}

			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return err
			}

			resp.Body.Close()
			fmt.Printf(". ")

			if resp.StatusCode == 200 {
				type Result struct {
					Result string `json:"result"`
				}

				res := Result{}

				err := json.Unmarshal(bytes.TrimSpace(body), &res)
				if err != nil {
					return err
				}

				fmt.Println(res.Result)
				return nil
			}
		case <-timeout:
			return errors.New("timeout")
		}
	}

	return nil
}

func PrintBox(msg string) {
	fmt.Printf("\n\n")
	fmt.Println("╔" + strings.Repeat("═", len(msg)+2) + "╗")
	fmt.Println("║" + strings.Repeat(" ", len(msg)+2) + "║")
	fmt.Println("║ " + msg + " ║")
	fmt.Println("║" + strings.Repeat(" ", len(msg)+2) + "║")
	fmt.Println("╚" + strings.Repeat("═", len(msg)+2) + "╝")
	fmt.Printf("\n")
}

func (k *Kd) GetConfig() error {
	var err error
	// we except to read all of them, if one fails it should be abort
	k.PublicKey, err = GetKey("public")
	if err != nil {
		return err
	}

	k.PrivateKey, err = GetKey("private")
	if err != nil {
		return err
	}
	return nil
}

func GetKey(key string) (string, error) {
	usr, err := user.Current()
	if err != nil {
		return "", err
	}

	var keyfile string
	switch key {
	case "public":
		keyfile = usr.HomeDir + "/.kd/koding.key.pub"
	case "private":
		keyfile = usr.HomeDir + "/.kd/koding.key"
	default:
		return "", fmt.Errorf("key is not recognized '%s'\n", key)
	}

	file, err := ioutil.ReadFile(keyfile)
	if err != nil {
		return "", err
	}

	return strings.TrimSpace(string(file)), nil
}
<commit_msg>register rcvr fixed<commit_after>package cli

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	uuid "github.com/nu7hatch/gouuid"
	"io/ioutil"
	"net/http"
	"os/user"
	"strings"
	"time"
)

// var authSite = "https://koding.com"

var authSite = "http://localhost:3020"

type Kd struct {
	ID         string
	PublicKey  string
	PrivateKey string
	Username   string
}

func NewKd() *Kd {
	id, _ := uuid.NewV4()
	return &Kd{
		ID: "machineID-" + id.String(),
	}
}

func (k *Kd) Definition() string {
	return "Registers the kite system to kontrol"
}

func (k *Kd) Exec() error {
	err := k.Register()
	if err != nil {
		return errors.New(fmt.Sprint("\nregister error: %s\n", err.Error()))
	}
	return nil
}

func (k *Kd) Register() error {
	err := k.GetConfig()
	if err != nil {
		return err
	}

	registerUrl := fmt.Sprintf("%s/-/auth/register/%s/%s", authSite, k.ID, k.PublicKey)
	checkUrl := fmt.Sprintf("%s/-/auth/check/%s", authSite, k.PublicKey)

	fmt.Printf("Please open the following url for authentication:\n\n")
	// PrintBox(registerUrl)
	fmt.Println(registerUrl)
	fmt.Printf("\nwaiting . ")

	// check the result every two seconds
	ticker := time.NewTicker(time.Second * 2).C

	// wait for three minutes, if not successfull abort it
	timeout := time.After(time.Minute * 3)

	for {
		select {
		case <-ticker:
			resp, err := http.Get(checkUrl)
			if err != nil {
				return err
			}

			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return err
			}

			resp.Body.Close()
			fmt.Printf(". ")

			if resp.StatusCode == 200 {
				type Result struct {
					Result string `json:"result"`
				}

				res := Result{}

				err := json.Unmarshal(bytes.TrimSpace(body), &res)
				if err != nil {
					return err
				}

				fmt.Println(res.Result)
				return nil
			}
		case <-timeout:
			return errors.New("timeout")
		}
	}

	return nil
}

func PrintBox(msg string) {
	fmt.Printf("\n\n")
	fmt.Println("╔" + strings.Repeat("═", len(msg)+2) + "╗")
	fmt.Println("║" + strings.Repeat(" ", len(msg)+2) + "║")
	fmt.Println("║ " + msg + " ║")
	fmt.Println("║" + strings.Repeat(" ", len(msg)+2) + "║")
	fmt.Println("╚" + strings.Repeat("═", len(msg)+2) + "╝")
	fmt.Printf("\n")
}

func (k *Kd) GetConfig() error {
	var err error
	// we except to read all of them, if one fails it should be abort
	k.PublicKey, err = GetKey("public")
	if err != nil {
		return err
	}

	k.PrivateKey, err = GetKey("private")
	if err != nil {
		return err
	}
	return nil
}

func GetKey(key string) (string, error) {
	usr, err := user.Current()
	if err != nil {
		return "", err
	}

	var keyfile string
	switch key {
	case "public":
		keyfile = usr.HomeDir + "/.kd/koding.key.pub"
	case "private":
		keyfile = usr.HomeDir + "/.kd/koding.key"
	default:
		return "", fmt.Errorf("key is not recognized '%s'\n", key)
	}

	file, err := ioutil.ReadFile(keyfile)
	if err != nil {
		return "", err
	}

	return strings.TrimSpace(string(file)), nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"

	"github.com/codegangsta/cli"
	"github.com/xenolf/lego/acme"
)

func checkFolder(path string) error {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return os.MkdirAll(path, 0700)
	}
	return nil
}

func setup(c *cli.Context) (*Configuration, *Account, *acme.Client) {
	err := checkFolder(c.GlobalString("path"))
	if err != nil {
		logger().Fatalf("Cound not check/create path: %v", err)
	}

	conf := NewConfiguration(c)
	if !c.GlobalIsSet("email") {
		logger().Fatal("You have to pass an account (email address) to the program using --email or -m")
	}

	//TODO: move to account struct? Currently MUST pass email.
	acc := NewAccount(c.GlobalString("email"), conf)

	client, err := acme.NewClient(c.GlobalString("server"), acc, conf.RsaBits(), conf.OptPort())
	if err != nil {
		logger().Fatal("Could not create client:", err)
	}

	return conf, acc, client
}

func saveCertRes(certRes acme.CertificateResource, conf *Configuration) {
	// We store the certificate, private key and metadata in different files
	// as web servers would not be able to work with a combined file.
	certOut := path.Join(conf.CertPath(), certRes.Domain+".crt")
	privOut := path.Join(conf.CertPath(), certRes.Domain+".key")
	metaOut := path.Join(conf.CertPath(), certRes.Domain+".json")

	err := ioutil.WriteFile(certOut, certRes.Certificate, 0600)
	if err != nil {
		logger().Printf("Unable to save Certificate for domain %s\n\t%v", certRes.Domain, err)
	}

	err = ioutil.WriteFile(privOut, certRes.PrivateKey, 0600)
	if err != nil {
		logger().Printf("Unable to save PrivateKey for domain %s\n\t%v", certRes.Domain, err)
	}

	jsonBytes, err := json.MarshalIndent(certRes, "", "\t")
	if err != nil {
		logger().Printf("Unable to marshal CertResource for domain %s\n\t%v", certRes.Domain, err)
	}

	err = ioutil.WriteFile(metaOut, jsonBytes, 0600)
	if err != nil {
		logger().Printf("Unable to save CertResource for domain %s\n\t%v", certRes.Domain, err)
	}
}

func run(c *cli.Context) {
	conf, acc, client := setup(c)
	if acc.Registration == nil {
		reg, err := client.Register()
		if err != nil {
			logger().Fatalf("Could not complete registration\n\t%v", err)
		}

		acc.Registration = reg
		acc.Save()

		logger().Print("!!!! HEADS UP !!!!")
		logger().Printf(`
		Your account credentials have been saved in your Let's Encrypt
		configuration directory at "%s".
		You should make a secure backup	of this folder now. This
		configuration directory will also contain certificates and
		private keys obtained from Let's Encrypt so making regular
		backups of this folder is ideal.`, conf.AccountPath(c.GlobalString("email")))

	}

	if acc.Registration.Body.Agreement == "" {
		reader := bufio.NewReader(os.Stdin)
		logger().Printf("Please review the TOS at %s", acc.Registration.TosURL)

		for {
			logger().Println("Do you accept the TOS? Y/n")
			text, err := reader.ReadString('\n')
			if err != nil {
				logger().Fatalf("Could not read from console -> %v", err)
			}

			text = strings.Trim(text, "\r\n")

			if text == "n" {
				logger().Fatal("You did not accept the TOS. Unable to proceed.")
			}

			if text == "Y" || text == "y" || text == "" {
				err = client.AgreeToTOS()
				if err != nil {
					logger().Fatalf("Could not agree to tos -> %v", err)
				}
				acc.Save()
				break
			}

			logger().Println("Your input was invalid. Please answer with one of Y/y, n or by pressing enter.")
		}
	}

	if !c.GlobalIsSet("domains") {
		logger().Fatal("Please specify --domains")
	}

	cert, failures := client.ObtainSANCertificate(c.GlobalStringSlice("domains"), true)
	if len(failures) > 0 {
		for k, v := range failures {
			logger().Printf("[%s] Could not obtain certificates\n\t%v", k, v)
		}
	}

	err := checkFolder(conf.CertPath())
	if err != nil {
		logger().Fatalf("Cound not check/create path: %v", err)
	}

	saveCertRes(cert, conf)
}

func revoke(c *cli.Context) {

	conf, _, client := setup(c)

	err := checkFolder(conf.CertPath())
	if err != nil {
		logger().Fatalf("Cound not check/create path: %v", err)
	}

	for _, domain := range c.GlobalStringSlice("domains") {
		logger().Printf("Trying to revoke certificate for domain %s", domain)

		certPath := path.Join(conf.CertPath(), domain+".crt")
		certBytes, err := ioutil.ReadFile(certPath)

		err = client.RevokeCertificate(certBytes)
		if err != nil {
			logger().Printf("Error while revoking the certificate for domain %s\n\t%v", domain, err)
		} else {
			logger().Print("Certificate was revoked.")
		}
	}
}

func renew(c *cli.Context) {
	conf, _, client := setup(c)

	for _, domain := range c.GlobalStringSlice("domains") {
		// load the cert resource from files.
		// We store the certificate, private key and metadata in different files
		// as web servers would not be able to work with a combined file.
		certPath := path.Join(conf.CertPath(), domain+".crt")
		privPath := path.Join(conf.CertPath(), domain+".key")
		metaPath := path.Join(conf.CertPath(), domain+".json")

		certBytes, err := ioutil.ReadFile(certPath)
		if err != nil {
			logger().Printf("Error while loading the certificate for domain %s\n\t%v", domain, err)
			return
		}

		if c.IsSet("days") {
			expTime, err := acme.GetPEMCertExpiration(certBytes)
			if err != nil {
				logger().Printf("Could not get Certification expiration for domain %s", domain)
			}

			if int(expTime.Sub(time.Now()).Hours() / 24.0) <= c.Int("days") {
				continue
			}
		}

		keyBytes, err := ioutil.ReadFile(privPath)
		if err != nil {
			logger().Printf("Error while loading the private key for domain %s\n\t%v", domain, err)
			return
		}

		metaBytes, err := ioutil.ReadFile(metaPath)
		if err != nil {
			logger().Printf("Error while loading the meta data for domain %s\n\t%v", domain, err)
			return
		}

		var certRes acme.CertificateResource
		err = json.Unmarshal(metaBytes, &certRes)
		if err != nil {
			logger().Printf("Error while marshalling the meta data for domain %s\n\t%v", domain, err)
			return
		}

		certRes.PrivateKey = keyBytes
		certRes.Certificate = certBytes

		newCert, err := client.RenewCertificate(certRes, true, true)
		if err != nil {
			logger().Printf("%v", err)
			return
		}

		saveCertRes(newCert, conf)
	}
}
<commit_msg>Make sure the CLI returns proper non-zero exit codes on errors.<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"

	"github.com/codegangsta/cli"
	"github.com/xenolf/lego/acme"
)

func checkFolder(path string) error {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return os.MkdirAll(path, 0700)
	}
	return nil
}

func setup(c *cli.Context) (*Configuration, *Account, *acme.Client) {
	err := checkFolder(c.GlobalString("path"))
	if err != nil {
		logger().Fatalf("Cound not check/create path: %s", err.Error())
	}

	conf := NewConfiguration(c)
	if !c.GlobalIsSet("email") {
		logger().Fatal("You have to pass an account (email address) to the program using --email or -m")
	}

	//TODO: move to account struct? Currently MUST pass email.
	acc := NewAccount(c.GlobalString("email"), conf)

	client, err := acme.NewClient(c.GlobalString("server"), acc, conf.RsaBits(), conf.OptPort())
	if err != nil {
		logger().Fatalf("Could not create client: %s", err.Error())
	}

	return conf, acc, client
}

func saveCertRes(certRes acme.CertificateResource, conf *Configuration) {
	// We store the certificate, private key and metadata in different files
	// as web servers would not be able to work with a combined file.
	certOut := path.Join(conf.CertPath(), certRes.Domain+".crt")
	privOut := path.Join(conf.CertPath(), certRes.Domain+".key")
	metaOut := path.Join(conf.CertPath(), certRes.Domain+".json")

	err := ioutil.WriteFile(certOut, certRes.Certificate, 0600)
	if err != nil {
		logger().Fatalf("Unable to save Certificate for domain %s\n\t%s", certRes.Domain, err.Error())
	}

	err = ioutil.WriteFile(privOut, certRes.PrivateKey, 0600)
	if err != nil {
		logger().Fatalf("Unable to save PrivateKey for domain %s\n\t%s", certRes.Domain, err.Error())
	}

	jsonBytes, err := json.MarshalIndent(certRes, "", "\t")
	if err != nil {
		logger().Fatalf("Unable to marshal CertResource for domain %s\n\t%s", certRes.Domain, err.Error())
	}

	err = ioutil.WriteFile(metaOut, jsonBytes, 0600)
	if err != nil {
		logger().Fatalf("Unable to save CertResource for domain %s\n\t%s", certRes.Domain, err.Error())
	}
}

func run(c *cli.Context) {
	conf, acc, client := setup(c)
	if acc.Registration == nil {
		reg, err := client.Register()
		if err != nil {
			logger().Fatalf("Could not complete registration\n\t%s", err.Error())
		}

		acc.Registration = reg
		acc.Save()

		logger().Print("!!!! HEADS UP !!!!")
		logger().Printf(`
		Your account credentials have been saved in your Let's Encrypt
		configuration directory at "%s".
		You should make a secure backup	of this folder now. This
		configuration directory will also contain certificates and
		private keys obtained from Let's Encrypt so making regular
		backups of this folder is ideal.`, conf.AccountPath(c.GlobalString("email")))

	}

	if acc.Registration.Body.Agreement == "" {
		reader := bufio.NewReader(os.Stdin)
		logger().Printf("Please review the TOS at %s", acc.Registration.TosURL)

		for {
			logger().Println("Do you accept the TOS? Y/n")
			text, err := reader.ReadString('\n')
			if err != nil {
				logger().Fatalf("Could not read from console -> %s", err.Error())
			}

			text = strings.Trim(text, "\r\n")

			if text == "n" {
				logger().Fatal("You did not accept the TOS. Unable to proceed.")
			}

			if text == "Y" || text == "y" || text == "" {
				err = client.AgreeToTOS()
				if err != nil {
					logger().Fatalf("Could not agree to tos -> %s", err)
				}
				acc.Save()
				break
			}

			logger().Println("Your input was invalid. Please answer with one of Y/y, n or by pressing enter.")
		}
	}

	if !c.GlobalIsSet("domains") {
		logger().Fatal("Please specify --domains")
	}

	cert, failures := client.ObtainSANCertificate(c.GlobalStringSlice("domains"), true)
	if len(failures) > 0 {
		for k, v := range failures {
			logger().Printf("[%s] Could not obtain certificates\n\t%s", k, v.Error())
		}

		// Make sure to return a non-zero exit code if ObtainSANCertificate
		// returned at least one error. Due to us not returning partial
		// certificate we can just exit here instead of at the end.
		os.Exit(1)
	}

	err := checkFolder(conf.CertPath())
	if err != nil {
		logger().Fatalf("Cound not check/create path: %s", err.Error())
	}

	saveCertRes(cert, conf)
}

func revoke(c *cli.Context) {

	conf, _, client := setup(c)

	err := checkFolder(conf.CertPath())
	if err != nil {
		logger().Fatalf("Cound not check/create path: %s", err.Error())
	}

	for _, domain := range c.GlobalStringSlice("domains") {
		logger().Printf("Trying to revoke certificate for domain %s", domain)

		certPath := path.Join(conf.CertPath(), domain+".crt")
		certBytes, err := ioutil.ReadFile(certPath)

		err = client.RevokeCertificate(certBytes)
		if err != nil {
			logger().Fatalf("Error while revoking the certificate for domain %s\n\t%s", domain, err.Error())
		} else {
			logger().Print("Certificate was revoked.")
		}
	}
}

func renew(c *cli.Context) {
	conf, _, client := setup(c)

	for _, domain := range c.GlobalStringSlice("domains") {
		// load the cert resource from files.
		// We store the certificate, private key and metadata in different files
		// as web servers would not be able to work with a combined file.
		certPath := path.Join(conf.CertPath(), domain+".crt")
		privPath := path.Join(conf.CertPath(), domain+".key")
		metaPath := path.Join(conf.CertPath(), domain+".json")

		certBytes, err := ioutil.ReadFile(certPath)
		if err != nil {
			logger().Fatalf("Error while loading the certificate for domain %s\n\t%s", domain, err.Error())
		}

		if c.IsSet("days") {
			expTime, err := acme.GetPEMCertExpiration(certBytes)
			if err != nil {
				logger().Printf("Could not get Certification expiration for domain %s", domain)
			}

			if int(expTime.Sub(time.Now()).Hours() / 24.0) <= c.Int("days") {
				continue
			}
		}

		keyBytes, err := ioutil.ReadFile(privPath)
		if err != nil {
			logger().Fatalf("Error while loading the private key for domain %s\n\t%s", domain, err.Error())
		}

		metaBytes, err := ioutil.ReadFile(metaPath)
		if err != nil {
			logger().Fatalf("Error while loading the meta data for domain %s\n\t%s", domain, err.Error())
		}

		var certRes acme.CertificateResource
		err = json.Unmarshal(metaBytes, &certRes)
		if err != nil {
			logger().Fatalf("Error while marshalling the meta data for domain %s\n\t%s", domain, err.Error())
		}

		certRes.PrivateKey = keyBytes
		certRes.Certificate = certBytes

		newCert, err := client.RenewCertificate(certRes, true, true)
		if err != nil {
			logger().Fatalf("%s", err.Error())
		}

		saveCertRes(newCert, conf)
	}
}
<|endoftext|>


<commit_before>package perceptron

type Perceptron struct {
	iteration int
}

func New(iteration byte) *Perceptron {
	p := &Perceptron{
		iteration: int(iteration),
	}

	return p
}

func (p *Perceptron) Learn(model Model, iter InstanceIter) error {
	// TODO: Implement this.
	return nil
}

type Learner interface {
	Learn(model Model, iter InstanceIter) error
}

type Model interface {
}

type Instance interface {
}

type InstanceIter interface {
}
<commit_msg>"(*Perceptron).Learn" apply "(*Perceptron).learnInstance" to all instances.<commit_after>package perceptron

type Perceptron struct {
	iteration int
}

func New(iteration byte) *Perceptron {
	p := &Perceptron{
		iteration: int(iteration),
	}

	return p
}

func (p *Perceptron) Learn(model Model, iter InstanceIter) error {
	for iter.HasNext() {
		instance := iter.Get()

		if err := p.learnInstance(model, instance); err != nil {
			return err
		}
	}

	if err := iter.Error(); err != nil {
		return err
	}

	return nil
}

func (p *Perceptron) learnInstance(model Model, instance Instance) error {
	// TODO: Implement this.
	return nil
}

type Learner interface {
	Learn(model Model, iter InstanceIter) error
}

type Model interface {
}

type Instance interface {
}

type InstanceIter interface {
	HasNext() bool
	Get() Instance
	Error() error
}
<|endoftext|>


<commit_before>package pes

import "errors"
import "io"
import "io/ioutil"
import "github.com/32bitkid/mpeg/util"
import "github.com/32bitkid/mpeg/ts"

const (
	StartCodePrefix = 0x000001
)

var (
	ErrStartCodePrefixNotFound = errors.New("start code prefix not found")
	ErrMarkerNotFound          = errors.New("marker not found")
	ErrInvalidStuffingByte     = errors.New("invalid stuffing byte")
)

type Packet struct {
	StreamID     uint32
	PacketLength uint32
	*Header
	Payload []byte
}

func ReadPacketFrom(br util.BitReader32) (packet *Packet, err error) {
	packet = new(Packet)
	err = packet.readFrom(br)
	return packet, err
}

func (packet *Packet) readFrom(reader util.BitReader32) error {

	var (
		val uint32
		err error
	)

	val, err = reader.Peek32(24)

	if err != nil {
		return err
	}

	if val != StartCodePrefix {
		return ErrStartCodePrefixNotFound
	}

	reader.Trash(24)

	packet.StreamID, err = reader.Read32(8)
	if err != nil {
		return err
	}

	packet.PacketLength, err = reader.Read32(16)
	if err != nil {
		return err
	}

	switch {
	case hasPESHeader(packet.StreamID):
		var headerLen uint32
		packet.Header, headerLen, err = ReadHeader(reader)
		if err != nil {
			return err
		}

		if packet.PacketLength > 0 {
			var payloadLen = int(packet.PacketLength - headerLen)
			packet.Payload = make([]byte, payloadLen)
			_, err = io.ReadAtLeast(reader, packet.Payload, payloadLen)
			if err != nil {
				return err
			}
		} else {
			// Read until end of buffer
			packet.Payload, err = ioutil.ReadAll(reader)
			if err != nil && err != ts.EOP {
				return err
			}
		}
	case packet.StreamID == padding_stream:
		payloadLen := int(packet.PacketLength)
		junk := make([]byte, payloadLen)
		_, err = io.ReadAtLeast(reader, junk, payloadLen)
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>adding todo to revisit later<commit_after>package pes

import "errors"
import "io"
import "io/ioutil"
import "github.com/32bitkid/mpeg/util"
import "github.com/32bitkid/mpeg/ts"

const (
	StartCodePrefix = 0x000001
)

var (
	ErrStartCodePrefixNotFound = errors.New("start code prefix not found")
	ErrMarkerNotFound          = errors.New("marker not found")
	ErrInvalidStuffingByte     = errors.New("invalid stuffing byte")
)

type Packet struct {
	StreamID     uint32
	PacketLength uint32
	*Header
	Payload []byte
}

func ReadPacketFrom(br util.BitReader32) (packet *Packet, err error) {
	packet = new(Packet)
	err = packet.readFrom(br)
	return packet, err
}

func (packet *Packet) readFrom(reader util.BitReader32) error {

	var (
		val uint32
		err error
	)

	val, err = reader.Peek32(24)

	if err != nil {
		return err
	}

	if val != StartCodePrefix {
		return ErrStartCodePrefixNotFound
	}

	reader.Trash(24)

	packet.StreamID, err = reader.Read32(8)
	if err != nil {
		return err
	}

	packet.PacketLength, err = reader.Read32(16)
	if err != nil {
		return err
	}

	switch {
	case hasPESHeader(packet.StreamID):
		var headerLen uint32
		packet.Header, headerLen, err = ReadHeader(reader)
		if err != nil {
			return err
		}

		if packet.PacketLength > 0 {
			var payloadLen = int(packet.PacketLength - headerLen)
			packet.Payload = make([]byte, payloadLen)
			_, err = io.ReadAtLeast(reader, packet.Payload, payloadLen)
			if err != nil {
				return err
			}
		} else {
			// Read until end of buffer
			// TODO dont create a new slice but read into a shared sized buffer
			packet.Payload, err = ioutil.ReadAll(reader)
			if err != nil && err != ts.EOP {
				return err
			}
		}
	case packet.StreamID == padding_stream:
		payloadLen := int(packet.PacketLength)
		junk := make([]byte, payloadLen)
		_, err = io.ReadAtLeast(reader, junk, payloadLen)
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 Hajime Hoshi
//
// 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 loop

import (
	"errors"
	"sync"
	"time"

	"github.com/hajimehoshi/ebiten/internal/graphics/opengl"
	"github.com/hajimehoshi/ebiten/internal/ui"
)

func CurrentFPS() float64 {
	return currentRunContext.getCurrentFPS()
}

func IsRunning() bool {
	if currentRunContext == nil {
		return false
	}
	return currentRunContext.isRunning()
}

type runContext struct {
	running        bool
	fps            int
	currentFPS     float64
	runningSlowly  bool
	frames         int
	lastUpdated    int64
	lastFPSUpdated int64
	m              sync.RWMutex
}

var currentRunContext *runContext

func (c *runContext) startRunning() {
	c.m.Lock()
	defer c.m.Unlock()
	c.running = true
}

func (c *runContext) isRunning() bool {
	c.m.RLock()
	defer c.m.RUnlock()
	return c.running
}

func (c *runContext) endRunning() {
	c.m.Lock()
	defer c.m.Unlock()
	c.running = false
}

func (c *runContext) getCurrentFPS() float64 {
	c.m.RLock()
	defer c.m.RUnlock()
	if !c.running {
		// TODO: Should panic here?
		return 0
	}
	return c.currentFPS
}

func (c *runContext) updateFPS(fps float64) {
	c.m.Lock()
	defer c.m.Unlock()
	c.currentFPS = fps
}

type GraphicsContext interface {
	SetSize(width, height int, scale float64) error
	UpdateAndDraw(context *opengl.Context, updateCount int) error
}

type regularTermination struct{}

func (e regularTermination) Error() string {
	return "regular termination"
}

func Run(g GraphicsContext, width, height int, scale float64, title string, fps int) (err error) {
	if currentRunContext != nil {
		return errors.New("loop: The game is already running")
	}
	currentRunContext = &runContext{
		fps: fps,
	}
	currentRunContext.startRunning()
	defer currentRunContext.endRunning()

	if err := ui.CurrentUI().Start(width, height, scale, title); err != nil {
		return err
	}
	// TODO: Use the error value
	defer func() {
		_ = ui.CurrentUI().Terminate()
	}()

	n := now()
	currentRunContext.lastUpdated = n
	currentRunContext.lastFPSUpdated = n

	if err := ui.CurrentUI().AnimationFrameLoop(func() error {
		return currentRunContext.update(g)
	}); err != nil {
		if err == (regularTermination{}) {
			return nil
		}
		return err
	}
	return nil
}

func (c *runContext) update(g GraphicsContext) error {
	for {
		e, err := ui.CurrentUI().Update()
		if err != nil {
			return err
		}
		switch e := e.(type) {
		case ui.NopEvent:
			return nil
		case ui.ScreenSizeEvent:
			if err := g.SetSize(e.Width, e.Height, e.ActualScale); err != nil {
				return err
			}
			e.Done <- struct{}{}
			continue
		case ui.CloseEvent:
			return regularTermination{}
		case ui.RenderEvent:
			if err := currentRunContext.render(g); err != nil {
				return err
			}
			e.Done <- struct{}{}
			return nil
		default:
			panic("not reach")
		}
	}
}

func (c *runContext) render(g GraphicsContext) error {
	fps := c.fps
	n := now()
	defer func() {
		// Calc the current FPS.
		if time.Second > time.Duration(n-c.lastFPSUpdated) {
			return
		}
		currentFPS := float64(c.frames) * float64(time.Second) / float64(n-c.lastFPSUpdated)
		c.updateFPS(currentFPS)
		c.lastFPSUpdated = n
		c.frames = 0
	}()

	// If lastUpdated is too old, we assume that screen is not shown.
	if 10*int64(time.Second)/int64(fps) < n-c.lastUpdated {
		c.lastUpdated = n
		return nil
	}

	// Note that generally t is a little different from 1/60[sec].
	t := n - c.lastUpdated
	tt := int(t * int64(fps) / int64(time.Second))

	// As t is not accurate 1/60[sec], errors are accumulated.
	// To make the FPS stable, set tt 1 if t is a little less than 1/60[sec].
	if tt == 0 && (int64(time.Second)/int64(fps)-int64(5*time.Millisecond)) < t {
		tt = 1
	}
	if err := g.UpdateAndDraw(ui.GLContext(), tt); err != nil {
		return err
	}
	c.lastUpdated += int64(tt) * int64(time.Second) / int64(fps)
	c.frames++
	return nil
}
<commit_msg>loop: Simplify update<commit_after>// Copyright 2016 Hajime Hoshi
//
// 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 loop

import (
	"errors"
	"sync"
	"time"

	"github.com/hajimehoshi/ebiten/internal/graphics/opengl"
	"github.com/hajimehoshi/ebiten/internal/ui"
)

func CurrentFPS() float64 {
	return currentRunContext.getCurrentFPS()
}

func IsRunning() bool {
	if currentRunContext == nil {
		return false
	}
	return currentRunContext.isRunning()
}

type runContext struct {
	running        bool
	fps            int
	currentFPS     float64
	runningSlowly  bool
	frames         int
	lastUpdated    int64
	lastFPSUpdated int64
	m              sync.RWMutex
}

var currentRunContext *runContext

func (c *runContext) startRunning() {
	c.m.Lock()
	defer c.m.Unlock()
	c.running = true
}

func (c *runContext) isRunning() bool {
	c.m.RLock()
	defer c.m.RUnlock()
	return c.running
}

func (c *runContext) endRunning() {
	c.m.Lock()
	defer c.m.Unlock()
	c.running = false
}

func (c *runContext) getCurrentFPS() float64 {
	c.m.RLock()
	defer c.m.RUnlock()
	if !c.running {
		// TODO: Should panic here?
		return 0
	}
	return c.currentFPS
}

func (c *runContext) updateFPS(fps float64) {
	c.m.Lock()
	defer c.m.Unlock()
	c.currentFPS = fps
}

type GraphicsContext interface {
	SetSize(width, height int, scale float64) error
	UpdateAndDraw(context *opengl.Context, updateCount int) error
}

type regularTermination struct{}

func (e regularTermination) Error() string {
	return "regular termination"
}

func Run(g GraphicsContext, width, height int, scale float64, title string, fps int) (err error) {
	if currentRunContext != nil {
		return errors.New("loop: The game is already running")
	}
	currentRunContext = &runContext{
		fps: fps,
	}
	currentRunContext.startRunning()
	defer currentRunContext.endRunning()

	if err := ui.CurrentUI().Start(width, height, scale, title); err != nil {
		return err
	}
	// TODO: Use the error value
	defer func() {
		_ = ui.CurrentUI().Terminate()
	}()

	n := now()
	currentRunContext.lastUpdated = n
	currentRunContext.lastFPSUpdated = n

	if err := ui.CurrentUI().AnimationFrameLoop(func() error {
		return currentRunContext.update(g)
	}); err != nil {
		if err == (regularTermination{}) {
			return nil
		}
		return err
	}
	return nil
}

func (c *runContext) update(g GraphicsContext) error {
	e, err := ui.CurrentUI().Update()
	if err != nil {
		return err
	}
	switch e := e.(type) {
	case ui.NopEvent:
	case ui.ScreenSizeEvent:
		if err := g.SetSize(e.Width, e.Height, e.ActualScale); err != nil {
			return err
		}
		e.Done <- struct{}{}
	case ui.CloseEvent:
		return regularTermination{}
	case ui.RenderEvent:
		if err := currentRunContext.render(g); err != nil {
			return err
		}
		e.Done <- struct{}{}
	default:
		panic("not reach")
	}
	return nil
}

func (c *runContext) render(g GraphicsContext) error {
	fps := c.fps
	n := now()
	defer func() {
		// Calc the current FPS.
		if time.Second > time.Duration(n-c.lastFPSUpdated) {
			return
		}
		currentFPS := float64(c.frames) * float64(time.Second) / float64(n-c.lastFPSUpdated)
		c.updateFPS(currentFPS)
		c.lastFPSUpdated = n
		c.frames = 0
	}()

	// If lastUpdated is too old, we assume that screen is not shown.
	if 10*int64(time.Second)/int64(fps) < n-c.lastUpdated {
		c.lastUpdated = n
		return nil
	}

	// Note that generally t is a little different from 1/60[sec].
	t := n - c.lastUpdated
	tt := int(t * int64(fps) / int64(time.Second))

	// As t is not accurate 1/60[sec], errors are accumulated.
	// To make the FPS stable, set tt 1 if t is a little less than 1/60[sec].
	if tt == 0 && (int64(time.Second)/int64(fps)-int64(5*time.Millisecond)) < t {
		tt = 1
	}
	if err := g.UpdateAndDraw(ui.GLContext(), tt); err != nil {
		return err
	}
	c.lastUpdated += int64(tt) * int64(time.Second) / int64(fps)
	c.frames++
	return nil
}
<|endoftext|>


<commit_before>// Package ipfix decodes IPFIX packets
//: ----------------------------------------------------------------------------
//: Copyright (C) 2017 Verizon.  All Rights Reserved.
//: All Rights Reserved
//:
//: file:    reader_test.go
//: details: unit testing for reader.go
//: author:  Mehrdad Arshad Rad
//: date:    03/22/2017
//:
//: 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 ipfix

import (
	"reflect"
	"testing"
)

func TestUint8(t *testing.T) {
	b := []byte{0x05, 0x11}

	r := NewReader(b)
	i, err := r.Uint8()
	if err != nil {
		t.Error("unexpected error happened, got", err)
	}

	if i != 5 {
		t.Error("expect read 5, got", i)
	}
}

func TestUint16(t *testing.T) {
	b := []byte{0x05, 0x11}

	r := NewReader(b)
	i, err := r.Uint16()
	if err != nil {
		t.Error("unexpected error happened, got", err)
	}

	if i != 1297 {
		t.Error("expect read 1297, got", i)
	}
}

func TestUint32(t *testing.T) {
	b := []byte{0x05, 0x11, 0x01, 0x16}

	r := NewReader(b)
	i, err := r.Uint32()
	if err != nil {
		t.Error("unexpected error happened, got", err)
	}

	if i != 85000470 {
		t.Error("expect read 85000470, got", i)
	}
}

func TestReadN(t *testing.T) {
	b := []byte{0x05, 0x11, 0x01, 0x16}

	r := NewReader(b)
	i, err := r.Read(2)
	if err != nil {
		t.Error("unexpected error happened, got", err)
	}

	if !reflect.DeepEqual(i, []byte{0x05, 0x11}) {
		t.Error("expect read [5 17], got", i)
	}
}
<commit_msg>add uint64 testing<commit_after>// Package ipfix decodes IPFIX packets
//: ----------------------------------------------------------------------------
//: Copyright (C) 2017 Verizon.  All Rights Reserved.
//: All Rights Reserved
//:
//: file:    reader_test.go
//: details: unit testing for reader.go
//: author:  Mehrdad Arshad Rad
//: date:    03/22/2017
//:
//: 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 ipfix

import (
	"reflect"
	"testing"
)

func TestUint8(t *testing.T) {
	b := []byte{0x05, 0x11}

	r := NewReader(b)
	i, err := r.Uint8()
	if err != nil {
		t.Error("unexpected error happened, got", err)
	}

	if i != 5 {
		t.Error("expect read 5, got", i)
	}
}

func TestUint16(t *testing.T) {
	b := []byte{0x05, 0x11}

	r := NewReader(b)
	i, err := r.Uint16()
	if err != nil {
		t.Error("unexpected error happened, got", err)
	}

	if i != 1297 {
		t.Error("expect read 1297, got", i)
	}
}

func TestUint32(t *testing.T) {
	b := []byte{0x05, 0x11, 0x01, 0x16}

	r := NewReader(b)
	i, err := r.Uint32()
	if err != nil {
		t.Error("unexpected error happened, got", err)
	}

	if i != 85000470 {
		t.Error("expect read 85000470, got", i)
	}
}

func TestUint64(t *testing.T) {
	b := []byte{0x05, 0x11, 0x01, 0x16, 0x05, 0x01, 0x21, 0x26}

	r := NewReader(b)
	i, err := r.Uint64()
	if err != nil {
		t.Error("unexpected error happened, got", err)
	}

	if i != 365074238878589222 {
		t.Error("expect read 365074238878589222, got", i)
	}
}

func TestReadN(t *testing.T) {
	b := []byte{0x05, 0x11, 0x01, 0x16}

	r := NewReader(b)
	i, err := r.Read(2)
	if err != nil {
		t.Error("unexpected error happened, got", err)
	}

	if !reflect.DeepEqual(i, []byte{0x05, 0x11}) {
		t.Error("expect read [5 17], got", i)
	}
}
<|endoftext|>


<commit_before>package iptables

import (
	"bufio"
	"bytes"
	"os/exec"
	"regexp"
	"strconv"
	"sync"
	"time"

	"github.com/deckarep/golang-set"
	"github.com/luizbafilho/fusis/config"
	"github.com/luizbafilho/fusis/net"
	"github.com/luizbafilho/fusis/state"
	"github.com/luizbafilho/fusis/types"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
)

var (
	ErrIptablesNotFound = errors.New("[iptables] Binary not found")
	ErrIptablesSnat     = errors.New("[iptables] Error when inserting SNAT rule")
	ErrIptablesRule     = errors.New("[iptables] Error adding rule")
)

// Defines iptables actions
var (
	ADD = "-A"
	DEL = "-D"
)

type Syncer interface {
	Sync(state state.State) error
}

type IptablesMngr struct {
	sync.Mutex
	config *config.BalancerConfig
	path   string
}

type SnatRule struct {
	vaddr    string
	vport    string
	toSource string
}

func New(config *config.BalancerConfig) (*IptablesMngr, error) {
	// look for iptables
	path, err := exec.LookPath("iptables")
	if err != nil {
		return nil, ErrIptablesNotFound
	}

	// create FUSIS iptables chain, ignore error in case FUSIS chain already exists
	_ = exec.Command(path, "-t", "nat", "-N", "FUSIS").Run()

	// if FUSIS chain is not present in POSTROUTING
	err = exec.Command(path, "-t", "nat", "-C", "POSTROUTING", "-j", "FUSIS").Run()
	if err != nil {
		// add FUSIS chain
		err = exec.Command(path, "-t", "nat", "-A", "POSTROUTING", "-j", "FUSIS").Run()
		if err != nil {
			return nil, ErrIptablesRule
		}
	}

	if err := enableContrack(); err != nil {
		return nil, err
	}

	return &IptablesMngr{
		config: config,
		path:   path,
	}, nil
}

// Sync syncs all iptables rules
func (i IptablesMngr) Sync(s state.State) error {
	start := time.Now()
	defer func() {
		log.Debugf("[iptables ] Sync took %v", time.Since(start))
	}()

	log.Debug("[iptables] Syncing")

	stateSet, err := i.getStateRulesSet(s)
	if err != nil {
		return err
	}

	kernelSet, err := i.getKernelRulesSet()
	if err != nil {
		return err
	}

	rulesToAdd := stateSet.Difference(kernelSet)
	rulesToRemove := kernelSet.Difference(stateSet)

	// Adding missing rules
	for r := range rulesToAdd.Iter() {
		rule := r.(SnatRule)
		err := i.addRule(rule)
		if err != nil {
			return err
		}
		log.Debugf("[iptables] Added rule: %#v", rule)
	}

	// Cleaning rules
	for r := range rulesToRemove.Iter() {
		rule := r.(SnatRule)
		err := i.removeRule(rule)
		if err != nil {
			return err
		}
		log.Debugf("[iptables] Removed rule: %#v", rule)
	}

	return nil
}

func (i IptablesMngr) getKernelRulesSet() (mapset.Set, error) {
	kernelRules, err := i.getSnatRules()
	if err != nil {
		return nil, err
	}

	kernelSet := mapset.NewSet()
	for _, r := range kernelRules {
		kernelSet.Add(r)
	}

	return kernelSet, nil

}

func (i IptablesMngr) getStateRulesSet(s state.State) (mapset.Set, error) {
	stateRules, err := i.convertServicesToSnatRules(i.getNatServices(s))
	if err != nil {
		return nil, err
	}

	stateSet := mapset.NewSet()
	for _, r := range stateRules {
		stateSet.Add(r)
	}

	return stateSet, nil
}

func (i IptablesMngr) convertServicesToSnatRules(svcs []types.Service) ([]SnatRule, error) {
	rules := []SnatRule{}
	for _, s := range svcs {

		r, err := i.serviceToSnatRule(s)
		if err != nil {
			return rules, err
		}

		rules = append(rules, *r)
	}

	return rules, nil
}

func (i IptablesMngr) getNatServices(s state.State) []types.Service {
	natServices := []types.Service{}

	for _, svc := range s.GetServices() {
		if svc.IsNat() {
			natServices = append(natServices, svc)
		}
	}

	return natServices
}

func (i IptablesMngr) serviceToSnatRule(svc types.Service) (*SnatRule, error) {
	privateIp, err := net.GetIpByInterface(i.config.Interfaces.Outbound)
	if err != nil {
		return nil, err
	}

	rule := &SnatRule{
		vaddr:    svc.Address,
		vport:    strconv.Itoa(int(svc.Port)),
		toSource: privateIp,
	}

	return rule, nil
}

func (i IptablesMngr) addRule(r SnatRule) error {
	return i.execIptablesCommand(ADD, r)
}

func (i IptablesMngr) removeRule(r SnatRule) error {
	return i.execIptablesCommand(DEL, r)
}

func (i IptablesMngr) execIptablesCommand(action string, r SnatRule) error {
	i.Lock()
	defer i.Unlock()
	cmd := exec.Command(i.path, "-t", "nat", action, "FUSIS", "-m", "ipvs", "--vaddr", r.vaddr+"/32", "--vport", r.vport, "-j", "SNAT", "--to-source", r.toSource)
	err := cmd.Run()
	if err != nil {
		return err
	}

	return nil
}

func (i IptablesMngr) getSnatRules() ([]SnatRule, error) {
	i.Lock()
	defer i.Unlock()
	out, err := exec.Command(i.path, "--list", "FUSIS", "-t", "nat").Output()
	if err != nil {
		log.Fatal("[iptables] Error executing iptables ", err)
	}

	r, _ := regexp.Compile(`vaddr\s([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})\svport\s(\d+)\sto:([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})`)
	scanner := bufio.NewScanner(bytes.NewReader(out))

	rules := []SnatRule{}
	for scanner.Scan() {
		matches := r.FindStringSubmatch(scanner.Text())
		if len(matches) == 0 {
			continue
		}

		rules = append(rules, SnatRule{
			vaddr:    matches[1],
			vport:    matches[2],
			toSource: matches[3],
		})
	}

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

	return rules, nil
}

func enableContrack() error {
	if err := net.SetSysctl("net.ipv4.vs.conntrack", "1"); err != nil {
		return errors.Wrap(err, "enabling Contrack failed")
	}

	return nil
}
<commit_msg>fixing iptables test<commit_after>package iptables

import (
	"bufio"
	"bytes"
	"os/exec"
	"regexp"
	"strconv"
	"sync"
	"time"

	"github.com/deckarep/golang-set"
	"github.com/luizbafilho/fusis/config"
	"github.com/luizbafilho/fusis/net"
	"github.com/luizbafilho/fusis/state"
	"github.com/luizbafilho/fusis/types"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
)

var (
	ErrIptablesNotFound = errors.New("[iptables] Binary not found")
	ErrIptablesSnat     = errors.New("[iptables] Error when inserting SNAT rule")
	ErrIptablesRule     = errors.New("[iptables] Error adding rule")
)

// Defines iptables actions
var (
	ADD = "-A"
	DEL = "-D"
)

type Syncer interface {
	Sync(state state.State) error
}

type IptablesMngr struct {
	sync.Mutex
	config *config.BalancerConfig
	path   string
}

type SnatRule struct {
	vaddr    string
	vport    string
	toSource string
}

func New(config *config.BalancerConfig) (*IptablesMngr, error) {
	// look for iptables
	path, err := exec.LookPath("iptables")
	if err != nil {
		return nil, ErrIptablesNotFound
	}

	// create FUSIS iptables chain, ignore error in case FUSIS chain already exists
	_ = exec.Command(path, "-t", "nat", "-N", "FUSIS").Run()

	// if FUSIS chain is not present in POSTROUTING
	err = exec.Command(path, "-t", "nat", "-C", "POSTROUTING", "-j", "FUSIS").Run()
	if err != nil {
		// add FUSIS chain
		err = exec.Command(path, "-t", "nat", "-A", "POSTROUTING", "-j", "FUSIS").Run()
		if err != nil {
			return nil, ErrIptablesRule
		}
	}

	if err := enableContrack(); err != nil {
		return nil, err
	}

	return &IptablesMngr{
		config: config,
		path:   path,
	}, nil
}

// Sync syncs all iptables rules
func (i IptablesMngr) Sync(s state.State) error {
	start := time.Now()
	defer func() {
		log.Debugf("[iptables ] Sync took %v", time.Since(start))
	}()

	log.Debug("[iptables] Syncing")

	stateSet, err := i.getStateRulesSet(s)
	if err != nil {
		return err
	}

	kernelSet, err := i.getKernelRulesSet()
	if err != nil {
		return err
	}

	rulesToAdd := stateSet.Difference(kernelSet)
	rulesToRemove := kernelSet.Difference(stateSet)

	// Adding missing rules
	for r := range rulesToAdd.Iter() {
		rule := r.(SnatRule)
		err := i.addRule(rule)
		if err != nil {
			return err
		}
		log.Debugf("[iptables] Added rule: %#v", rule)
	}

	// Cleaning rules
	for r := range rulesToRemove.Iter() {
		rule := r.(SnatRule)
		err := i.removeRule(rule)
		if err != nil {
			return err
		}
		log.Debugf("[iptables] Removed rule: %#v", rule)
	}

	return nil
}

func (i IptablesMngr) getKernelRulesSet() (mapset.Set, error) {
	kernelRules, err := i.getSnatRules()
	if err != nil {
		return nil, err
	}

	kernelSet := mapset.NewSet()
	for _, r := range kernelRules {
		kernelSet.Add(r)
	}

	return kernelSet, nil

}

func (i IptablesMngr) getStateRulesSet(s state.State) (mapset.Set, error) {
	stateRules, err := i.convertServicesToSnatRules(i.getNatServices(s))
	if err != nil {
		return nil, err
	}

	stateSet := mapset.NewSet()
	for _, r := range stateRules {
		stateSet.Add(r)
	}

	return stateSet, nil
}

func (i IptablesMngr) convertServicesToSnatRules(svcs []types.Service) ([]SnatRule, error) {
	rules := []SnatRule{}
	for _, s := range svcs {

		r, err := i.serviceToSnatRule(s)
		if err != nil {
			return rules, err
		}

		rules = append(rules, *r)
	}

	return rules, nil
}

func (i IptablesMngr) getNatServices(s state.State) []types.Service {
	natServices := []types.Service{}

	for _, svc := range s.GetServices() {
		if svc.IsNat() {
			natServices = append(natServices, svc)
		}
	}

	return natServices
}

func (i IptablesMngr) serviceToSnatRule(svc types.Service) (*SnatRule, error) {
	privateIp, err := net.GetIpByInterface(i.config.Interfaces.Outbound)
	if err != nil {
		return nil, err
	}

	rule := &SnatRule{
		vaddr:    svc.Address,
		vport:    strconv.Itoa(int(svc.Port)),
		toSource: privateIp,
	}

	return rule, nil
}

func (i IptablesMngr) addRule(r SnatRule) error {
	return i.execIptablesCommand(ADD, r)
}

func (i IptablesMngr) removeRule(r SnatRule) error {
	return i.execIptablesCommand(DEL, r)
}

func (i IptablesMngr) execIptablesCommand(action string, r SnatRule) error {
	i.Lock()
	defer i.Unlock()
	cmd := exec.Command(i.path, "-t", "nat", action, "FUSIS", "-m", "ipvs", "--vaddr", r.vaddr+"/32", "--vport", r.vport, "-j", "SNAT", "--to-source", r.toSource)
	err := cmd.Run()
	if err != nil {
		return err
	}

	return nil
}

func (i IptablesMngr) getSnatRules() ([]SnatRule, error) {
	i.Lock()
	defer i.Unlock()
	out, err := exec.Command(i.path, "--list", "FUSIS", "-t", "nat", "-n").Output()
	if err != nil {
		log.Fatal("[iptables] Error executing iptables ", err)
	}

	r, _ := regexp.Compile(`vaddr\s([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})\svport\s(\d+)\sto:([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})`)
	scanner := bufio.NewScanner(bytes.NewReader(out))

	rules := []SnatRule{}
	for scanner.Scan() {
		matches := r.FindStringSubmatch(scanner.Text())
		if len(matches) == 0 {
			continue
		}

		rules = append(rules, SnatRule{
			vaddr:    matches[1],
			vport:    matches[2],
			toSource: matches[3],
		})
	}

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

	return rules, nil
}

func enableContrack() error {
	if err := net.SetSysctl("net.ipv4.vs.conntrack", "1"); err != nil {
		return errors.Wrap(err, "enabling Contrack failed")
	}

	return nil
}
<|endoftext|>


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

// +build darwin dragonfly freebsd linux netbsd openbsd

package ipv4

import (
	"os"
	"syscall"
	"unsafe"

	"golang.org/x/net/internal/iana"
)

func setControlMessage(fd int, opt *rawOpt, cf ControlFlags, on bool) error {
	opt.Lock()
	defer opt.Unlock()
	if cf&FlagTTL != 0 {
		if err := setInt(fd, &sockOpts[ssoReceiveTTL], boolint(on)); err != nil {
			return err
		}
		if on {
			opt.set(FlagTTL)
		} else {
			opt.clear(FlagTTL)
		}
	}
	if sockOpts[ssoPacketInfo].name > 0 {
		if cf&(FlagSrc|FlagDst|FlagInterface) != 0 {
			if err := setInt(fd, &sockOpts[ssoPacketInfo], boolint(on)); err != nil {
				return err
			}
			if on {
				opt.set(cf & (FlagSrc | FlagDst | FlagInterface))
			} else {
				opt.clear(cf & (FlagSrc | FlagDst | FlagInterface))
			}
		}
	} else {
		if cf&FlagDst != 0 {
			if err := setInt(fd, &sockOpts[ssoReceiveDst], boolint(on)); err != nil {
				return err
			}
			if on {
				opt.set(FlagDst)
			} else {
				opt.clear(FlagDst)
			}
		}
		if cf&FlagInterface != 0 {
			if err := setInt(fd, &sockOpts[ssoReceiveInterface], boolint(on)); err != nil {
				return err
			}
			if on {
				opt.set(FlagInterface)
			} else {
				opt.clear(FlagInterface)
			}
		}
	}
	return nil
}

func newControlMessage(opt *rawOpt) (oob []byte) {
	opt.RLock()
	var l int
	if opt.isset(FlagTTL) {
		l += syscall.CmsgSpace(ctlOpts[ctlTTL].length)
	}
	if ctlOpts[ctlPacketInfo].name > 0 {
		if opt.isset(FlagSrc | FlagDst | FlagInterface) {
			l += syscall.CmsgSpace(ctlOpts[ctlPacketInfo].length)
		}
	} else {
		if opt.isset(FlagDst) {
			l += syscall.CmsgSpace(ctlOpts[ctlDst].length)
		}
		if opt.isset(FlagInterface) {
			l += syscall.CmsgSpace(ctlOpts[ctlInterface].length)
		}
	}
	if l > 0 {
		oob = make([]byte, l)
		b := oob
		if opt.isset(FlagTTL) {
			b = ctlOpts[ctlTTL].marshal(b, nil)
		}
		if ctlOpts[ctlPacketInfo].name > 0 {
			if opt.isset(FlagSrc | FlagDst | FlagInterface) {
				b = ctlOpts[ctlPacketInfo].marshal(b, nil)
			}
		} else {
			if opt.isset(FlagDst) {
				b = ctlOpts[ctlDst].marshal(b, nil)
			}
			if opt.isset(FlagInterface) {
				b = ctlOpts[ctlInterface].marshal(b, nil)
			}
		}
	}
	opt.RUnlock()
	return
}

func parseControlMessage(b []byte) (*ControlMessage, error) {
	if len(b) == 0 {
		return nil, nil
	}
	cmsgs, err := syscall.ParseSocketControlMessage(b)
	if err != nil {
		return nil, os.NewSyscallError("parse socket control message", err)
	}
	cm := &ControlMessage{}
	for _, m := range cmsgs {
		if m.Header.Level != iana.ProtocolIP {
			continue
		}
		switch int(m.Header.Type) {
		case ctlOpts[ctlTTL].name:
			ctlOpts[ctlTTL].parse(cm, m.Data[:])
		case ctlOpts[ctlDst].name:
			ctlOpts[ctlDst].parse(cm, m.Data[:])
		case ctlOpts[ctlInterface].name:
			ctlOpts[ctlInterface].parse(cm, m.Data[:])
		case ctlOpts[ctlPacketInfo].name:
			ctlOpts[ctlPacketInfo].parse(cm, m.Data[:])
		}
	}
	return cm, nil
}

func marshalControlMessage(cm *ControlMessage) (oob []byte) {
	if cm == nil {
		return nil
	}
	var l int
	if ctlOpts[ctlPacketInfo].name > 0 {
		if cm.Src.To4() != nil || cm.IfIndex > 0 {
			l += syscall.CmsgSpace(ctlOpts[ctlPacketInfo].length)
		}
	}
	if l > 0 {
		oob = make([]byte, l)
		b := oob
		if ctlOpts[ctlPacketInfo].name > 0 {
			if cm.Src.To4() != nil || cm.IfIndex > 0 {
				b = ctlOpts[ctlPacketInfo].marshal(b, cm)
			}
		}
	}
	return
}

func marshalTTL(b []byte, cm *ControlMessage) []byte {
	m := (*syscall.Cmsghdr)(unsafe.Pointer(&b[0]))
	m.Level = iana.ProtocolIP
	m.Type = sysIP_RECVTTL
	m.SetLen(syscall.CmsgLen(1))
	return b[syscall.CmsgSpace(1):]
}

func parseTTL(cm *ControlMessage, b []byte) {
	cm.TTL = int(*(*byte)(unsafe.Pointer(&b[:1][0])))
}
<commit_msg>x/net/ipv4: add missing ancillary data socket options check<commit_after>// Copyright 2012 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build darwin dragonfly freebsd linux netbsd openbsd

package ipv4

import (
	"os"
	"syscall"
	"unsafe"

	"golang.org/x/net/internal/iana"
)

func setControlMessage(fd int, opt *rawOpt, cf ControlFlags, on bool) error {
	opt.Lock()
	defer opt.Unlock()
	if cf&FlagTTL != 0 && sockOpts[ssoReceiveTTL].name > 0 {
		if err := setInt(fd, &sockOpts[ssoReceiveTTL], boolint(on)); err != nil {
			return err
		}
		if on {
			opt.set(FlagTTL)
		} else {
			opt.clear(FlagTTL)
		}
	}
	if sockOpts[ssoPacketInfo].name > 0 {
		if cf&(FlagSrc|FlagDst|FlagInterface) != 0 {
			if err := setInt(fd, &sockOpts[ssoPacketInfo], boolint(on)); err != nil {
				return err
			}
			if on {
				opt.set(cf & (FlagSrc | FlagDst | FlagInterface))
			} else {
				opt.clear(cf & (FlagSrc | FlagDst | FlagInterface))
			}
		}
	} else {
		if cf&FlagDst != 0 && sockOpts[ssoReceiveDst].name > 0 {
			if err := setInt(fd, &sockOpts[ssoReceiveDst], boolint(on)); err != nil {
				return err
			}
			if on {
				opt.set(FlagDst)
			} else {
				opt.clear(FlagDst)
			}
		}
		if cf&FlagInterface != 0 && sockOpts[ssoReceiveInterface].name > 0 {
			if err := setInt(fd, &sockOpts[ssoReceiveInterface], boolint(on)); err != nil {
				return err
			}
			if on {
				opt.set(FlagInterface)
			} else {
				opt.clear(FlagInterface)
			}
		}
	}
	return nil
}

func newControlMessage(opt *rawOpt) (oob []byte) {
	opt.RLock()
	var l int
	if opt.isset(FlagTTL) && ctlOpts[ctlTTL].name > 0 {
		l += syscall.CmsgSpace(ctlOpts[ctlTTL].length)
	}
	if ctlOpts[ctlPacketInfo].name > 0 {
		if opt.isset(FlagSrc | FlagDst | FlagInterface) {
			l += syscall.CmsgSpace(ctlOpts[ctlPacketInfo].length)
		}
	} else {
		if opt.isset(FlagDst) && ctlOpts[ctlDst].name > 0 {
			l += syscall.CmsgSpace(ctlOpts[ctlDst].length)
		}
		if opt.isset(FlagInterface) && ctlOpts[ctlInterface].name > 0 {
			l += syscall.CmsgSpace(ctlOpts[ctlInterface].length)
		}
	}
	if l > 0 {
		oob = make([]byte, l)
		b := oob
		if opt.isset(FlagTTL) && ctlOpts[ctlTTL].name > 0 {
			b = ctlOpts[ctlTTL].marshal(b, nil)
		}
		if ctlOpts[ctlPacketInfo].name > 0 {
			if opt.isset(FlagSrc | FlagDst | FlagInterface) {
				b = ctlOpts[ctlPacketInfo].marshal(b, nil)
			}
		} else {
			if opt.isset(FlagDst) && ctlOpts[ctlDst].name > 0 {
				b = ctlOpts[ctlDst].marshal(b, nil)
			}
			if opt.isset(FlagInterface) && ctlOpts[ctlInterface].name > 0 {
				b = ctlOpts[ctlInterface].marshal(b, nil)
			}
		}
	}
	opt.RUnlock()
	return
}

func parseControlMessage(b []byte) (*ControlMessage, error) {
	if len(b) == 0 {
		return nil, nil
	}
	cmsgs, err := syscall.ParseSocketControlMessage(b)
	if err != nil {
		return nil, os.NewSyscallError("parse socket control message", err)
	}
	cm := &ControlMessage{}
	for _, m := range cmsgs {
		if m.Header.Level != iana.ProtocolIP {
			continue
		}
		switch int(m.Header.Type) {
		case ctlOpts[ctlTTL].name:
			ctlOpts[ctlTTL].parse(cm, m.Data[:])
		case ctlOpts[ctlDst].name:
			ctlOpts[ctlDst].parse(cm, m.Data[:])
		case ctlOpts[ctlInterface].name:
			ctlOpts[ctlInterface].parse(cm, m.Data[:])
		case ctlOpts[ctlPacketInfo].name:
			ctlOpts[ctlPacketInfo].parse(cm, m.Data[:])
		}
	}
	return cm, nil
}

func marshalControlMessage(cm *ControlMessage) (oob []byte) {
	if cm == nil {
		return nil
	}
	var l int
	pktinfo := false
	if ctlOpts[ctlPacketInfo].name > 0 && (cm.Src.To4() != nil || cm.IfIndex > 0) {
		pktinfo = true
		l += syscall.CmsgSpace(ctlOpts[ctlPacketInfo].length)
	}
	if l > 0 {
		oob = make([]byte, l)
		b := oob
		if pktinfo {
			b = ctlOpts[ctlPacketInfo].marshal(b, cm)
		}
	}
	return
}

func marshalTTL(b []byte, cm *ControlMessage) []byte {
	m := (*syscall.Cmsghdr)(unsafe.Pointer(&b[0]))
	m.Level = iana.ProtocolIP
	m.Type = sysIP_RECVTTL
	m.SetLen(syscall.CmsgLen(1))
	return b[syscall.CmsgSpace(1):]
}

func parseTTL(cm *ControlMessage, b []byte) {
	cm.TTL = int(*(*byte)(unsafe.Pointer(&b[:1][0])))
}
<|endoftext|>