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

package main

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"regexp"
	"strings"
	"sync"
)

type Route struct {
	match   *regexp.Regexp
	backend string
}

// Config stores the TLS routing configuration.
type Config struct {
	mu     sync.Mutex
	routes []Route
}

func dnsRegex(s string) (*regexp.Regexp, error) {
	return regexp.Compile(s)
}

func (c *Config) Match(hostname string) string {
	c.mu.Lock()
	defer c.mu.Unlock()
	for _, r := range c.routes {
		if r.match.MatchString(hostname) {
			return r.backend
		}
	}
	return ""
}

func (c *Config) Read(r io.Reader) error {
	var routes []Route

	s := bufio.NewScanner(r)
	for s.Scan() {
		fs := strings.Fields(s.Text())
		switch len(fs) {
		case 0:
			continue
		case 1:
			return fmt.Errorf("invalid %q on a line by itself", s.Text())
		case 2:
			re, err := dnsRegex(fs[0])
			if err != nil {
				return err
			}
			routes = append(routes, Route{re, fs[1]})
		default:
			// TODO: multiple backends?
			return fmt.Errorf("too many fields on line: %q", s.Text())
		}
	}
	if err := s.Err(); err != nil {
		return err
	}

	c.mu.Lock()
	defer c.mu.Unlock()
	c.routes = routes
	return nil
}

func (c *Config) ReadFile(path string) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	return c.Read(f)
}
<commit_msg>Add DNS name support to config<commit_after>// Copyright 2016 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"regexp"
	"strings"
	"sync"
)

type Route struct {
	match   *regexp.Regexp
	backend string
}

// Config stores the TLS routing configuration.
type Config struct {
	mu     sync.Mutex
	routes []Route
}

func dnsRegex(s string) (*regexp.Regexp, error) {
	if len(s) >= 2 && s[0] == '/' && s[len(s)-1] == '/' {
		return regexp.Compile(s[1 : len(s)-1])
	}

	var b []string
	for _, f := range strings.Split(s, ".") {
		switch f {
		case "*":
			b = append(b, `[^.]+`)
		case "":
			return nil, fmt.Errorf("DNS name %q has empty label", s)
		default:
			b = append(b, regexp.QuoteMeta(f))
		}
	}
	return regexp.Compile(strings.Join(b, `\.`))
}

func (c *Config) Match(hostname string) string {
	c.mu.Lock()
	defer c.mu.Unlock()
	for _, r := range c.routes {
		if r.match.MatchString(hostname) {
			return r.backend
		}
	}
	return ""
}

func (c *Config) Read(r io.Reader) error {
	var routes []Route

	s := bufio.NewScanner(r)
	for s.Scan() {
		fs := strings.Fields(s.Text())
		switch len(fs) {
		case 0:
			continue
		case 1:
			return fmt.Errorf("invalid %q on a line by itself", s.Text())
		case 2:
			re, err := dnsRegex(fs[0])
			if err != nil {
				return err
			}
			routes = append(routes, Route{re, fs[1]})
		default:
			// TODO: multiple backends?
			return fmt.Errorf("too many fields on line: %q", s.Text())
		}
	}
	if err := s.Err(); err != nil {
		return err
	}

	c.mu.Lock()
	defer c.mu.Unlock()
	c.routes = routes
	return nil
}

func (c *Config) ReadFile(path string) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	return c.Read(f)
}
<|endoftext|>


<commit_before>/*
 * MumbleDJ
 * By Matthieu Grieger
 * bot/track_test.go
 * Copyright (c) 2016 Matthieu Grieger (MIT License)
 */

package bot

import (
	"testing"
	"time"

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

type TrackTestSuite struct {
	suite.Suite
	Track Track
}

func (suite *TrackTestSuite) SetupTest() {
	duration, _ := time.ParseDuration("1s")
	suite.Track = Track{
		ID:           "id",
		Title:        "title",
		Author:       "author",
		Submitter:    "submitter",
		Service:      "service",
		Filename:     "filename",
		ThumbnailURL: "thumbnailurl",
		Duration:     duration,
		Playlist:     new(Playlist),
	}
}

func (suite *TrackTestSuite) TestGetID() {
	suite.Equal("id", suite.Track.GetID())
}

func (suite *TrackTestSuite) TestGetTitle() {
	suite.Equal("title", suite.Track.GetTitle())
}

func (suite *TrackTestSuite) TestGetAuthor() {
	suite.Equal("author", suite.Track.GetAuthor())
}

func (suite *TrackTestSuite) TestGetSubmitter() {
	suite.Equal("submitter", suite.Track.GetSubmitter())
}

func (suite *TrackTestSuite) TestGetService() {
	suite.Equal("service", suite.Track.GetService())
}

func (suite *TrackTestSuite) TestGetFilenameWhenExists() {
	result := suite.Track.GetFilename()

	suite.Equal("filename", result)
}

func (suite *TrackTestSuite) TestGetFilenameWhenNotExists() {
	suite.Track.Filename = ""

	result := suite.Track.GetFilename()

	suite.Equal("", result)
}

func (suite *TrackTestSuite) TestGetThumbnailURLWhenExists() {
	result := suite.Track.GetThumbnailURL()

	suite.Equal("thumbnailurl", result)
}

func (suite *TrackTestSuite) TestGetThumbnailURLWhenNotExists() {
	suite.Track.ThumbnailURL = ""

	result := suite.Track.GetThumbnailURL()

	suite.Equal("", result)
}

func (suite *TrackTestSuite) TestGetDuration() {
	duration, _ := time.ParseDuration("1s")

	suite.Equal(duration, suite.Track.GetDuration())
}

func (suite *TrackTestSuite) TestGetPlaylistWhenExists() {
	result := suite.Track.GetPlaylist()

	suite.NotNil(result)
}

func (suite *TrackTestSuite) TestGetPlaylistWhenNotExists() {
	suite.Track.Playlist = nil

	result := suite.Track.GetPlaylist()

	suite.Nil(result)
}

func TestTrackTestSuite(t *testing.T) {
	suite.Run(t, new(TrackTestSuite))
}
<commit_msg>Implemented rest of track tests<commit_after>/*
 * MumbleDJ
 * By Matthieu Grieger
 * bot/track_test.go
 * Copyright (c) 2016 Matthieu Grieger (MIT License)
 */

package bot

import (
	"testing"
	"time"

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

type TrackTestSuite struct {
	suite.Suite
	Track Track
}

func (suite *TrackTestSuite) SetupTest() {
	duration, _ := time.ParseDuration("1s")
	offset, _ := time.ParseDuration("2ms")
	suite.Track = Track{
		ID:             "id",
		URL:            "url",
		Title:          "title",
		Author:         "author",
		AuthorURL:      "author_url",
		Submitter:      "submitter",
		Service:        "service",
		Filename:       "filename",
		ThumbnailURL:   "thumbnailurl",
		Duration:       duration,
		PlaybackOffset: offset,
		Playlist:       new(Playlist),
	}
}

func (suite *TrackTestSuite) TestGetID() {
	suite.Equal("id", suite.Track.GetID())
}

func (suite *TrackTestSuite) TestGetURL() {
	suite.Equal("url", suite.Track.GetURL())
}

func (suite *TrackTestSuite) TestGetTitle() {
	suite.Equal("title", suite.Track.GetTitle())
}

func (suite *TrackTestSuite) TestGetAuthor() {
	suite.Equal("author", suite.Track.GetAuthor())
}

func (suite *TrackTestSuite) TestGetAuthorURL() {
	suite.Equal("author_url", suite.Track.GetAuthorURL())
}

func (suite *TrackTestSuite) TestGetSubmitter() {
	suite.Equal("submitter", suite.Track.GetSubmitter())
}

func (suite *TrackTestSuite) TestGetService() {
	suite.Equal("service", suite.Track.GetService())
}

func (suite *TrackTestSuite) TestGetFilenameWhenExists() {
	result := suite.Track.GetFilename()

	suite.Equal("filename", result)
}

func (suite *TrackTestSuite) TestGetFilenameWhenNotExists() {
	suite.Track.Filename = ""

	result := suite.Track.GetFilename()

	suite.Equal("", result)
}

func (suite *TrackTestSuite) TestGetThumbnailURLWhenExists() {
	result := suite.Track.GetThumbnailURL()

	suite.Equal("thumbnailurl", result)
}

func (suite *TrackTestSuite) TestGetThumbnailURLWhenNotExists() {
	suite.Track.ThumbnailURL = ""

	result := suite.Track.GetThumbnailURL()

	suite.Equal("", result)
}

func (suite *TrackTestSuite) TestGetDuration() {
	duration, _ := time.ParseDuration("1s")

	suite.Equal(duration, suite.Track.GetDuration())
}

func (suite *TrackTestSuite) TestGetPlaybackOffset() {
	duration, _ := time.ParseDuration("2ms")

	suite.Equal(duration, suite.Track.GetPlaybackOffset())
}

func (suite *TrackTestSuite) TestGetPlaylistWhenExists() {
	result := suite.Track.GetPlaylist()

	suite.NotNil(result)
}

func (suite *TrackTestSuite) TestGetPlaylistWhenNotExists() {
	suite.Track.Playlist = nil

	result := suite.Track.GetPlaylist()

	suite.Nil(result)
}

func TestTrackTestSuite(t *testing.T) {
	suite.Run(t, new(TrackTestSuite))
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"math"
	"os"
	"strconv"
)

type config struct {
	Listen ListenConfig

	ClientConnInfo string

	StartupParameters map[string]string
	Databases VirtualDatabaseConfiguration

	Prometheus PrometheusConfig
}

var Config = config{
	// These are the defaults

	Listen:	ListenConfig{6433, "localhost", true},

	ClientConnInfo: "host=localhost port=5432 sslmode=disable",

	StartupParameters: nil,
	Databases: nil,

	Prometheus: PrometheusConfig{
		Enabled: false,
		Listen: ListenConfig{},
	},
}

func readIntValue(dst *int, val interface{}, option string) error {
	var err error

	switch val := val.(type) {
	case float64:
		if math.Trunc(val) == val {
			*dst = int(val)
		} else {
			err = fmt.Errorf("input must be an integer")
		}
	case string:
		*dst, err = strconv.Atoi(val)
	default:
		err = fmt.Errorf("input must be an integer")
	}
	if err != nil {
		return fmt.Errorf("invalid value for option %q: %s", option, err.Error())
	}
	return nil
}

func readTextValue(dst *string, val interface{}, option string) error {
	var err error

	switch val := val.(type) {
	case string:
		*dst = val
		err = nil
	default:
		err = fmt.Errorf("input must be a text string")
	}
	if err != nil {
		return fmt.Errorf("invalid value for option %q: %s", option, err.Error())
	}
	return nil
}

func readBooleanValue(dst *bool, val interface{}, option string) error {
	var err error

	switch val := val.(type) {
	case bool:
		*dst = val
		err = nil
	default:
		err = fmt.Errorf("input must be a boolean")
	}
	if err != nil {
		return fmt.Errorf("invalid value for option %q: %s", option, err.Error())
	}
	return nil
}

func readListenSection(c *ListenConfig, val interface{}, option string) error {
	data, ok := val.(map[string]interface{})
	if !ok {
		return fmt.Errorf(`section %q must be a JSON object`, option)
	}
	for key, value := range data {
		var err error

		switch key {
		case "port":
			err = readIntValue(&c.Port, value, option + ".port")
		case "host":
			err = readTextValue(&c.Host, value, option + ".host")
		case "keepalive":
			err = readBooleanValue(&c.KeepAlive, value, option + ".keepalive")
		default:
			err = fmt.Errorf("unrecognized configuration option %q", option+"."+key)
		}
		if err != nil {
			return err
		}
	}
	return nil
}

func readConnectSection(c *config, val interface{}) error {
	data, ok := val.(string)
	if !ok {
		return fmt.Errorf(`section "connect" must be a connection string`)
	}
	c.ClientConnInfo = data
	return nil
}

func readStartupParameterSection(c *config, val interface{}) error {
	data, ok := val.(map[string]interface{})
	if !ok {
		return fmt.Errorf(`section "startup_parameters" must be a set of key-value pairs`)
	}
	c.StartupParameters = make(map[string]string)
	for k, v := range data {
		vs, ok := v.(string)
		if !ok {
			return fmt.Errorf(`all startup parameters must be strings`)
		}
		c.StartupParameters[k] = vs
	}
	return nil
}

func readAuthSection(c *AuthConfig, val interface{}, option string) error {
	data, ok := val.(map[string]interface{})
	if !ok {
		return fmt.Errorf(`section %q must be a JSON object`, option)
	}

	for key, value := range data {
		var err error

		switch key {
		case "method":
			err = readTextValue(&c.method, value, option+".method")
		case "user":
			err = readTextValue(&c.user, value, option+".user")
		case "password":
			err = readTextValue(&c.password, value, option+".password")
		default:
			err = fmt.Errorf("unrecognized configuration option %q", option+"."+key)
		}
		if err != nil {
			return err
		}
	}

	switch c.method {
	case "md5":
	case "trust":
	default:
		return fmt.Errorf("unrecognized authentication method %q in %q", c.method, option)
	}

	return nil
}

func readDatabaseSection(c *config, val interface{}) error {
	array, ok := val.([]interface{})
	if !ok {
		return fmt.Errorf(`section "databases" must be a JSON array`)
	}

	for dbindex, el := range array {
		data, ok := el.(map[string]interface{})
		if !ok {
			return fmt.Errorf(`elements within the "databases" array must be JSON objects`)
		}

		option := fmt.Sprintf("databases[%d]", dbindex)
		var db virtualDatabase

		for key, value := range data {
			var err error

			switch key {
			case "name":
				err = readTextValue(&db.name, value, option+".name")
			case "auth":
				err = readAuthSection(&db.auth, value, option+".auth")
			default:
				err = fmt.Errorf("unrecognized configuration option %q", option+"."+key)
			}
			if err != nil {
				return err
			}
		}

		for _, pedb := range c.Databases {
			if pedb.name == db.name {
				return fmt.Errorf("database name %q is not unique", db.name)
			}
		}
		c.Databases = append(c.Databases, db)
	}

	return nil
}

func readPrometheusSection(c *config, val interface{}) error {
	data, ok := val.(map[string]interface{})
	if !ok {
		return fmt.Errorf(`section "prometheus" must be a JSON object`)
	}
	for key, value := range data {
		var err error

		switch key {
		case "listen":
			err = readListenSection(&c.Prometheus.Listen, value, "prometheus.listen")
			c.Prometheus.Enabled = true
		default:
			err = fmt.Errorf("unrecognized configuration option %q", "prometheus." + key)
		}
		if err != nil {
			return err
		}
	}
	return nil
}

func readConfigFile(filename string) error {
	var ci interface{}

	fh, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer fh.Close()
	dec := json.NewDecoder(fh)
	err = dec.Decode(&ci)
	if err != nil {
		return err
	}

	sections, ok := ci.(map[string]interface{})
	if !ok {
		return fmt.Errorf("configuration must be a JSON object")
	}

	for key, value := range sections {
		var err error

		switch key {
		case "listen":
			err = readListenSection(&Config.Listen, value, "listen")
		case "connect":
			err = readConnectSection(&Config, value)
		case "startup_parameters":
			err = readStartupParameterSection(&Config, value)
		case "databases":
			err = readDatabaseSection(&Config, value)
		case "prometheus":
			err = readPrometheusSection(&Config, value)
		default:
			err = fmt.Errorf("unrecognized configuration section %q", key)
		}
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Enable Prometheus if the section exists<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"math"
	"os"
	"strconv"
)

type config struct {
	Listen ListenConfig

	ClientConnInfo string

	StartupParameters map[string]string
	Databases VirtualDatabaseConfiguration

	Prometheus PrometheusConfig
}

var Config = config{
	// These are the defaults

	Listen:	ListenConfig{6433, "localhost", true},

	ClientConnInfo: "host=localhost port=5432 sslmode=disable",

	StartupParameters: nil,
	Databases: nil,

	Prometheus: PrometheusConfig{
		Enabled: false,
		Listen: ListenConfig{},
	},
}

func readIntValue(dst *int, val interface{}, option string) error {
	var err error

	switch val := val.(type) {
	case float64:
		if math.Trunc(val) == val {
			*dst = int(val)
		} else {
			err = fmt.Errorf("input must be an integer")
		}
	case string:
		*dst, err = strconv.Atoi(val)
	default:
		err = fmt.Errorf("input must be an integer")
	}
	if err != nil {
		return fmt.Errorf("invalid value for option %q: %s", option, err.Error())
	}
	return nil
}

func readTextValue(dst *string, val interface{}, option string) error {
	var err error

	switch val := val.(type) {
	case string:
		*dst = val
		err = nil
	default:
		err = fmt.Errorf("input must be a text string")
	}
	if err != nil {
		return fmt.Errorf("invalid value for option %q: %s", option, err.Error())
	}
	return nil
}

func readBooleanValue(dst *bool, val interface{}, option string) error {
	var err error

	switch val := val.(type) {
	case bool:
		*dst = val
		err = nil
	default:
		err = fmt.Errorf("input must be a boolean")
	}
	if err != nil {
		return fmt.Errorf("invalid value for option %q: %s", option, err.Error())
	}
	return nil
}

func readListenSection(c *ListenConfig, val interface{}, option string) error {
	data, ok := val.(map[string]interface{})
	if !ok {
		return fmt.Errorf(`section %q must be a JSON object`, option)
	}
	for key, value := range data {
		var err error

		switch key {
		case "port":
			err = readIntValue(&c.Port, value, option + ".port")
		case "host":
			err = readTextValue(&c.Host, value, option + ".host")
		case "keepalive":
			err = readBooleanValue(&c.KeepAlive, value, option + ".keepalive")
		default:
			err = fmt.Errorf("unrecognized configuration option %q", option+"."+key)
		}
		if err != nil {
			return err
		}
	}
	return nil
}

func readConnectSection(c *config, val interface{}) error {
	data, ok := val.(string)
	if !ok {
		return fmt.Errorf(`section "connect" must be a connection string`)
	}
	c.ClientConnInfo = data
	return nil
}

func readStartupParameterSection(c *config, val interface{}) error {
	data, ok := val.(map[string]interface{})
	if !ok {
		return fmt.Errorf(`section "startup_parameters" must be a set of key-value pairs`)
	}
	c.StartupParameters = make(map[string]string)
	for k, v := range data {
		vs, ok := v.(string)
		if !ok {
			return fmt.Errorf(`all startup parameters must be strings`)
		}
		c.StartupParameters[k] = vs
	}
	return nil
}

func readAuthSection(c *AuthConfig, val interface{}, option string) error {
	data, ok := val.(map[string]interface{})
	if !ok {
		return fmt.Errorf(`section %q must be a JSON object`, option)
	}

	for key, value := range data {
		var err error

		switch key {
		case "method":
			err = readTextValue(&c.method, value, option+".method")
		case "user":
			err = readTextValue(&c.user, value, option+".user")
		case "password":
			err = readTextValue(&c.password, value, option+".password")
		default:
			err = fmt.Errorf("unrecognized configuration option %q", option+"."+key)
		}
		if err != nil {
			return err
		}
	}

	switch c.method {
	case "md5":
	case "trust":
	default:
		return fmt.Errorf("unrecognized authentication method %q in %q", c.method, option)
	}

	return nil
}

func readDatabaseSection(c *config, val interface{}) error {
	array, ok := val.([]interface{})
	if !ok {
		return fmt.Errorf(`section "databases" must be a JSON array`)
	}

	for dbindex, el := range array {
		data, ok := el.(map[string]interface{})
		if !ok {
			return fmt.Errorf(`elements within the "databases" array must be JSON objects`)
		}

		option := fmt.Sprintf("databases[%d]", dbindex)
		var db virtualDatabase

		for key, value := range data {
			var err error

			switch key {
			case "name":
				err = readTextValue(&db.name, value, option+".name")
			case "auth":
				err = readAuthSection(&db.auth, value, option+".auth")
			default:
				err = fmt.Errorf("unrecognized configuration option %q", option+"."+key)
			}
			if err != nil {
				return err
			}
		}

		for _, pedb := range c.Databases {
			if pedb.name == db.name {
				return fmt.Errorf("database name %q is not unique", db.name)
			}
		}
		c.Databases = append(c.Databases, db)
	}

	return nil
}

func readPrometheusSection(c *config, val interface{}) error {
	data, ok := val.(map[string]interface{})
	if !ok {
		return fmt.Errorf(`section "prometheus" must be a JSON object`)
	}
	for key, value := range data {
		var err error

		switch key {
		case "listen":
			err = readListenSection(&c.Prometheus.Listen, value, "prometheus.listen")
		default:
			err = fmt.Errorf("unrecognized configuration option %q", "prometheus." + key)
		}
		if err != nil {
			return err
		}
	}
	c.Prometheus.Enabled = true
	return nil
}

func readConfigFile(filename string) error {
	var ci interface{}

	fh, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer fh.Close()
	dec := json.NewDecoder(fh)
	err = dec.Decode(&ci)
	if err != nil {
		return err
	}

	sections, ok := ci.(map[string]interface{})
	if !ok {
		return fmt.Errorf("configuration must be a JSON object")
	}

	for key, value := range sections {
		var err error

		switch key {
		case "listen":
			err = readListenSection(&Config.Listen, value, "listen")
		case "connect":
			err = readConnectSection(&Config, value)
		case "startup_parameters":
			err = readStartupParameterSection(&Config, value)
		case "databases":
			err = readDatabaseSection(&Config, value)
		case "prometheus":
			err = readPrometheusSection(&Config, value)
		default:
			err = fmt.Errorf("unrecognized configuration section %q", key)
		}
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package brightbox

import (
	"crypto/sha1"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"hash/crc32"
	"net/url"
	"os"
	"regexp"
	"strings"

	brightbox "github.com/brightbox/gobrightbox"
	"github.com/gophercloud/gophercloud"
	"github.com/gorhill/cronexpr"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
)

// const (
// 	maxPort = 65535
// 	minPort = 1
// )

var (
	serverRegexp           = regexp.MustCompile("^srv-.....$")
	serverGroupRegexp      = regexp.MustCompile("^grp-.....$")
	databaseTypeRegexp     = regexp.MustCompile("^dbt-.....$")
	databaseServerRegexp   = regexp.MustCompile("^dbs-.....$")
	databaseSnapshotRegexp = regexp.MustCompile("^dbi-.....$")
	loadBalancerRegexp     = regexp.MustCompile("^lba-.....$")
	zoneRegexp             = regexp.MustCompile("^(zon-.....$|gb1s?-[ab])$")
	serverTypeRegexp       = regexp.MustCompile("^typ-.....$")
	firewallPolicyRegexp   = regexp.MustCompile("^fwp-.....$")
	firewallRuleRegexp     = regexp.MustCompile("^fwr-.....$")
	interfaceRegexp        = regexp.MustCompile("^int-.....$")
	imageRegexp            = regexp.MustCompile("^img-.....$")
	dnsNameRegexp          = regexp.MustCompile("^(?:[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?.)+[a-z0-9][a-z0-9-]{0,61}[a-z0-9]$")
	unreadable             = map[string]bool{
		"deleted": true,
		"failed":  true,
	}
)

func hashString(
	v interface{},
) string {
	switch v.(type) {
	case string:
		return userDataHashSum(v.(string))
	default:
		return ""
	}
}

func userDataHashSum(userData string) string {
	// Check whether the userData is not Base64 encoded.
	// Always calculate hash of base64 decoded value since we
	// check against double-encoding when setting it
	v, base64DecodeError := base64Decode(userData)
	if base64DecodeError != nil {
		v = userData
	}
	hash := sha1.Sum([]byte(v))
	return hex.EncodeToString(hash[:])
}

func assignMap(d *schema.ResourceData, target **map[string]interface{}, index string) {
	if d.HasChange(index) {
		if attr, ok := d.GetOk(index); ok {
			temp := attr.(map[string]interface{})
			*target = &temp
		} else {
			temp := make(map[string]interface{})
			*target = &temp
		}
	}
}

func assignString(d *schema.ResourceData, target **string, index string) {
	if d.HasChange(index) {
		if *target == nil {
			var temp string
			*target = &temp
		}
		if attr, ok := d.GetOk(index); ok {
			**target = attr.(string)
		}
	}
}

func assignStringSet(d *schema.ResourceData, target *[]string, index string) {
	if d.HasChange(index) {
		*target = sliceFromStringSet(d, index)
	}
}

func sliceFromStringSet(d *schema.ResourceData, index string) []string {
	configured := d.Get(index).(*schema.Set).List()
	slice := make([]string, len(configured))
	for i, data := range configured {
		slice[i] = data.(string)
	}
	return slice
}

func flattenStringSlice(list []string) []interface{} {
	temp := make([]interface{}, len(list))
	for i, v := range list {
		temp[i] = v
	}
	return temp
}

func assignInt(d *schema.ResourceData, target **int, index string) {
	if d.HasChange(index) {
		var temp int
		if attr, ok := d.GetOk(index); ok {
			temp = attr.(int)
		}
		*target = &temp
	}
}

func assignBool(d *schema.ResourceData, target **bool, index string) {
	if d.HasChange(index) {
		var temp bool
		if attr, ok := d.GetOk(index); ok {
			temp = attr.(bool)
		}
		*target = &temp
	}
}

func setPrimaryCloudIP(d *schema.ResourceData, cloudIP *brightbox.CloudIP) {
	d.Set("ipv4_address", cloudIP.PublicIP)
	d.Set("public_hostname", cloudIP.Fqdn)
}

// Base64Encode encodes data if the input isn't already encoded
// using base64.StdEncoding.EncodeToString. If the input is already base64
// encoded, return the original input unchanged.
func base64Encode(data string) string {
	// Check whether the data is already Base64 encoded; don't double-encode
	if isBase64Encoded(data) {
		return data
	}
	// data has not been encoded encode and return
	return base64.StdEncoding.EncodeToString([]byte(data))
}

func isBase64Encoded(data string) bool {
	_, err := base64Decode(data)
	return err == nil
}

func base64Decode(data string) (string, error) {
	result, err := base64.StdEncoding.DecodeString(data)
	return string(result), err
}

func stringValidateFunc(v interface{}, name string, failureTest func(string) bool, formatString string) (warns []string, errors []error) {
	value := v.(string)
	if failureTest(value) {
		errors = append(errors, fmt.Errorf(formatString, name))
	}
	return
}

// ValidateCronString checks if the string is a valid cron layout
// An empty string is acceptable.
func ValidateCronString(v interface{}, name string) (warns []string, errors []error) {
	cronstr := v.(string)
	if cronstr == "" {
		return
	}
	if _, err := cronexpr.Parse(cronstr); err != nil {
		errors = append(errors, fmt.Errorf("%q not a valid Cron: %s", name, err))
	}
	return
}

func http1Keys(v interface{}, name string) (warns []string, errors []error) {
	mapValue, ok := v.(map[string]interface{})
	if !ok {
		errors = append(errors, fmt.Errorf("expected type of %s to be a Map", name))
		return
	}
	for k := range mapValue {
		if !validToken(k) {
			errors = append(errors, fmt.Errorf("Metadata key %s is an invalid token. Should be all lower case, with no underscores", k))

		}
	}
	return
}

// Token has to be lower case
func validToken(tok string) bool {
	for i := 0; i < len(tok); i++ {
		if !validHeaderFieldByte(tok[i]) {
			return false
		}
	}
	return true
}

// validHeaderFieldByte reports whether b is a valid byte in a header
// field name. RFC 7230 says:
//   header-field   = field-name ":" OWS field-value OWS
//   field-name     = token
//   tchar = "!" / "#" / "$" / "%" / "&" / "'" / "*" / "+" / "-" / "." /
//           "^" / "_" / "`" / "|" / "~" / DIGIT / ALPHA
//   token = 1*tchar
//
// Underscore isn't valid. Needs to be a hyphen as Swift silently
// converts otherwise.
func validHeaderFieldByte(b byte) bool {
	return int(b) < len(isTokenTable) && isTokenTable[b]
}

var isTokenTable = [127]bool{
	'!':  true,
	'#':  true,
	'$':  true,
	'%':  true,
	'&':  true,
	'\'': true,
	'*':  true,
	'+':  true,
	'-':  true,
	'.':  true,
	'0':  true,
	'1':  true,
	'2':  true,
	'3':  true,
	'4':  true,
	'5':  true,
	'6':  true,
	'7':  true,
	'8':  true,
	'9':  true,
	'A':  false,
	'B':  false,
	'C':  false,
	'D':  false,
	'E':  false,
	'F':  false,
	'G':  false,
	'H':  false,
	'I':  false,
	'J':  false,
	'K':  false,
	'L':  false,
	'M':  false,
	'N':  false,
	'O':  false,
	'P':  false,
	'Q':  false,
	'R':  false,
	'S':  false,
	'T':  false,
	'U':  false,
	'W':  false,
	'V':  false,
	'X':  false,
	'Y':  false,
	'Z':  false,
	'^':  true,
	'_':  false,
	'`':  true,
	'a':  true,
	'b':  true,
	'c':  true,
	'd':  true,
	'e':  true,
	'f':  true,
	'g':  true,
	'h':  true,
	'i':  true,
	'j':  true,
	'k':  true,
	'l':  true,
	'm':  true,
	'n':  true,
	'o':  true,
	'p':  true,
	'q':  true,
	'r':  true,
	's':  true,
	't':  true,
	'u':  true,
	'v':  true,
	'w':  true,
	'x':  true,
	'y':  true,
	'z':  true,
	'|':  true,
	'~':  true,
}

func escapedString(attr interface{}) string {
	return url.PathEscape(attr.(string))
}

func escapedStringList(source []string) []string {
	dest := make([]string, len(source))
	for i, v := range source {
		dest[i] = escapedString(v)
	}
	return dest
}

func escapedStringMetadata(metadata interface{}) map[string]string {
	source := metadata.(map[string]interface{})
	dest := make(map[string]string, len(source))
	for k, v := range source {
		dest[strings.ToLower(k)] = url.PathEscape(v.(string))
	}
	return dest
}

func removedMetadataKeys(old interface{}, new interface{}) []string {
	oldMap := old.(map[string]interface{})
	newMap := new.(map[string]interface{})
	result := make([]string, 0, len(oldMap))
	for key := range oldMap {
		if newMap[key] == nil {
			result = append(result, strings.ToLower(key))
		}
	}
	return result
}

// CheckDeleted checks the error to see if it's a 404 (Not Found) and, if so,
// sets the resource ID to the empty string instead of throwing an error.
func CheckDeleted(d *schema.ResourceData, err error, msg string) error {
	if _, ok := err.(gophercloud.ErrDefault404); ok {
		d.SetId("")
		return nil
	}

	return fmt.Errorf("%s %s: %s", msg, d.Id(), err)
}

// getEnvVarWithDefault retrieves the value of the environment variable
// named by the key. If the variable is not present, return the default
//value instead.
func getenvWithDefault(key string, defaultValue string) string {
	if val, exists := os.LookupEnv(key); exists {
		return val
	}
	return defaultValue
}

// set the lock state of a resource based upon a boolean
func setLockState(client *brightbox.Client, isLocked bool, resource interface{}) error {
	if isLocked {
		return client.LockResource(resource)
	}
	return client.UnLockResource(resource)
}

// strSliceContains checks if a given string is contained in a slice
// When anybody asks why Go needs generics, here you go.
func strSliceContains(haystack []string, needle string) bool {
	for _, s := range haystack {
		if s == needle {
			return true
		}
	}
	return false
}

// Check a JSON object is correct
func validateJSONObject(v interface{}, k string) ([]string, []error) {
	if v == nil || v.(string) == "" {
		return nil, []error{fmt.Errorf("%q value must not be empty", k)}
	}

	var j map[string]interface{}
	s := v.(string)

	err := json.Unmarshal([]byte(s), &j)
	if err != nil {
		return nil, []error{fmt.Errorf("%q must be a JSON object: %s", k, err)}
	}

	return nil, nil
}

func diffSuppressJSONObject(k, old, new string, d *schema.ResourceData) bool {
	if strSliceContains([]string{"{}", ""}, old) &&
		strSliceContains([]string{"{}", ""}, new) {
		return true
	}
	return false
}

// HashcodeString hashes a string to a unique hashcode.
//
// crc32 returns a uint32, but for our use we need
// and non negative integer. Here we cast to an integer
// and invert it if the result is negative.
func HashcodeString(s string) int {
	v := int(crc32.ChecksumIEEE([]byte(s)))
	if v >= 0 {
		return v
	}
	if -v >= 0 {
		return -v
	}
	// v == MinInt
	return 0
}

// StringIsValidFirewallTarget checks whether a string would
// pass the Iptables validation as a valid source or destination.
func stringIsValidFirewallTarget() schema.SchemaValidateFunc {
	return validation.Any(
		validation.StringInSlice([]string{"any"}, false),
		validation.StringMatch(serverRegexp, "must be a valid server ID"),
		validation.StringMatch(serverGroupRegexp, "must be a valid server Group ID"),
		validation.IsCIDR,
		validation.IsIPAddress,
	)
}

// Get a list of server IDs from a list of servers
func serverIDListFromNodes(
	nodes []brightbox.Server,
) []string {
	nodeIds := make([]string, len(nodes))
	for i, node := range nodes {
		nodeIds[i] = node.Id
	}
	return nodeIds
}
<commit_msg>Allow load balancer IDs in firewall rules<commit_after>package brightbox

import (
	"crypto/sha1"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"hash/crc32"
	"net/url"
	"os"
	"regexp"
	"strings"

	brightbox "github.com/brightbox/gobrightbox"
	"github.com/gophercloud/gophercloud"
	"github.com/gorhill/cronexpr"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
)

// const (
// 	maxPort = 65535
// 	minPort = 1
// )

var (
	serverRegexp           = regexp.MustCompile("^srv-.....$")
	serverGroupRegexp      = regexp.MustCompile("^grp-.....$")
	databaseTypeRegexp     = regexp.MustCompile("^dbt-.....$")
	databaseServerRegexp   = regexp.MustCompile("^dbs-.....$")
	databaseSnapshotRegexp = regexp.MustCompile("^dbi-.....$")
	loadBalancerRegexp     = regexp.MustCompile("^lba-.....$")
	zoneRegexp             = regexp.MustCompile("^(zon-.....$|gb1s?-[ab])$")
	serverTypeRegexp       = regexp.MustCompile("^typ-.....$")
	firewallPolicyRegexp   = regexp.MustCompile("^fwp-.....$")
	firewallRuleRegexp     = regexp.MustCompile("^fwr-.....$")
	interfaceRegexp        = regexp.MustCompile("^int-.....$")
	imageRegexp            = regexp.MustCompile("^img-.....$")
	dnsNameRegexp          = regexp.MustCompile("^(?:[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?.)+[a-z0-9][a-z0-9-]{0,61}[a-z0-9]$")
	unreadable             = map[string]bool{
		"deleted": true,
		"failed":  true,
	}
)

func hashString(
	v interface{},
) string {
	switch v.(type) {
	case string:
		return userDataHashSum(v.(string))
	default:
		return ""
	}
}

func userDataHashSum(userData string) string {
	// Check whether the userData is not Base64 encoded.
	// Always calculate hash of base64 decoded value since we
	// check against double-encoding when setting it
	v, base64DecodeError := base64Decode(userData)
	if base64DecodeError != nil {
		v = userData
	}
	hash := sha1.Sum([]byte(v))
	return hex.EncodeToString(hash[:])
}

func assignMap(d *schema.ResourceData, target **map[string]interface{}, index string) {
	if d.HasChange(index) {
		if attr, ok := d.GetOk(index); ok {
			temp := attr.(map[string]interface{})
			*target = &temp
		} else {
			temp := make(map[string]interface{})
			*target = &temp
		}
	}
}

func assignString(d *schema.ResourceData, target **string, index string) {
	if d.HasChange(index) {
		if *target == nil {
			var temp string
			*target = &temp
		}
		if attr, ok := d.GetOk(index); ok {
			**target = attr.(string)
		}
	}
}

func assignStringSet(d *schema.ResourceData, target *[]string, index string) {
	if d.HasChange(index) {
		*target = sliceFromStringSet(d, index)
	}
}

func sliceFromStringSet(d *schema.ResourceData, index string) []string {
	configured := d.Get(index).(*schema.Set).List()
	slice := make([]string, len(configured))
	for i, data := range configured {
		slice[i] = data.(string)
	}
	return slice
}

func flattenStringSlice(list []string) []interface{} {
	temp := make([]interface{}, len(list))
	for i, v := range list {
		temp[i] = v
	}
	return temp
}

func assignInt(d *schema.ResourceData, target **int, index string) {
	if d.HasChange(index) {
		var temp int
		if attr, ok := d.GetOk(index); ok {
			temp = attr.(int)
		}
		*target = &temp
	}
}

func assignBool(d *schema.ResourceData, target **bool, index string) {
	if d.HasChange(index) {
		var temp bool
		if attr, ok := d.GetOk(index); ok {
			temp = attr.(bool)
		}
		*target = &temp
	}
}

func setPrimaryCloudIP(d *schema.ResourceData, cloudIP *brightbox.CloudIP) {
	d.Set("ipv4_address", cloudIP.PublicIP)
	d.Set("public_hostname", cloudIP.Fqdn)
}

// Base64Encode encodes data if the input isn't already encoded
// using base64.StdEncoding.EncodeToString. If the input is already base64
// encoded, return the original input unchanged.
func base64Encode(data string) string {
	// Check whether the data is already Base64 encoded; don't double-encode
	if isBase64Encoded(data) {
		return data
	}
	// data has not been encoded encode and return
	return base64.StdEncoding.EncodeToString([]byte(data))
}

func isBase64Encoded(data string) bool {
	_, err := base64Decode(data)
	return err == nil
}

func base64Decode(data string) (string, error) {
	result, err := base64.StdEncoding.DecodeString(data)
	return string(result), err
}

func stringValidateFunc(v interface{}, name string, failureTest func(string) bool, formatString string) (warns []string, errors []error) {
	value := v.(string)
	if failureTest(value) {
		errors = append(errors, fmt.Errorf(formatString, name))
	}
	return
}

// ValidateCronString checks if the string is a valid cron layout
// An empty string is acceptable.
func ValidateCronString(v interface{}, name string) (warns []string, errors []error) {
	cronstr := v.(string)
	if cronstr == "" {
		return
	}
	if _, err := cronexpr.Parse(cronstr); err != nil {
		errors = append(errors, fmt.Errorf("%q not a valid Cron: %s", name, err))
	}
	return
}

func http1Keys(v interface{}, name string) (warns []string, errors []error) {
	mapValue, ok := v.(map[string]interface{})
	if !ok {
		errors = append(errors, fmt.Errorf("expected type of %s to be a Map", name))
		return
	}
	for k := range mapValue {
		if !validToken(k) {
			errors = append(errors, fmt.Errorf("Metadata key %s is an invalid token. Should be all lower case, with no underscores", k))

		}
	}
	return
}

// Token has to be lower case
func validToken(tok string) bool {
	for i := 0; i < len(tok); i++ {
		if !validHeaderFieldByte(tok[i]) {
			return false
		}
	}
	return true
}

// validHeaderFieldByte reports whether b is a valid byte in a header
// field name. RFC 7230 says:
//   header-field   = field-name ":" OWS field-value OWS
//   field-name     = token
//   tchar = "!" / "#" / "$" / "%" / "&" / "'" / "*" / "+" / "-" / "." /
//           "^" / "_" / "`" / "|" / "~" / DIGIT / ALPHA
//   token = 1*tchar
//
// Underscore isn't valid. Needs to be a hyphen as Swift silently
// converts otherwise.
func validHeaderFieldByte(b byte) bool {
	return int(b) < len(isTokenTable) && isTokenTable[b]
}

var isTokenTable = [127]bool{
	'!':  true,
	'#':  true,
	'$':  true,
	'%':  true,
	'&':  true,
	'\'': true,
	'*':  true,
	'+':  true,
	'-':  true,
	'.':  true,
	'0':  true,
	'1':  true,
	'2':  true,
	'3':  true,
	'4':  true,
	'5':  true,
	'6':  true,
	'7':  true,
	'8':  true,
	'9':  true,
	'A':  false,
	'B':  false,
	'C':  false,
	'D':  false,
	'E':  false,
	'F':  false,
	'G':  false,
	'H':  false,
	'I':  false,
	'J':  false,
	'K':  false,
	'L':  false,
	'M':  false,
	'N':  false,
	'O':  false,
	'P':  false,
	'Q':  false,
	'R':  false,
	'S':  false,
	'T':  false,
	'U':  false,
	'W':  false,
	'V':  false,
	'X':  false,
	'Y':  false,
	'Z':  false,
	'^':  true,
	'_':  false,
	'`':  true,
	'a':  true,
	'b':  true,
	'c':  true,
	'd':  true,
	'e':  true,
	'f':  true,
	'g':  true,
	'h':  true,
	'i':  true,
	'j':  true,
	'k':  true,
	'l':  true,
	'm':  true,
	'n':  true,
	'o':  true,
	'p':  true,
	'q':  true,
	'r':  true,
	's':  true,
	't':  true,
	'u':  true,
	'v':  true,
	'w':  true,
	'x':  true,
	'y':  true,
	'z':  true,
	'|':  true,
	'~':  true,
}

func escapedString(attr interface{}) string {
	return url.PathEscape(attr.(string))
}

func escapedStringList(source []string) []string {
	dest := make([]string, len(source))
	for i, v := range source {
		dest[i] = escapedString(v)
	}
	return dest
}

func escapedStringMetadata(metadata interface{}) map[string]string {
	source := metadata.(map[string]interface{})
	dest := make(map[string]string, len(source))
	for k, v := range source {
		dest[strings.ToLower(k)] = url.PathEscape(v.(string))
	}
	return dest
}

func removedMetadataKeys(old interface{}, new interface{}) []string {
	oldMap := old.(map[string]interface{})
	newMap := new.(map[string]interface{})
	result := make([]string, 0, len(oldMap))
	for key := range oldMap {
		if newMap[key] == nil {
			result = append(result, strings.ToLower(key))
		}
	}
	return result
}

// CheckDeleted checks the error to see if it's a 404 (Not Found) and, if so,
// sets the resource ID to the empty string instead of throwing an error.
func CheckDeleted(d *schema.ResourceData, err error, msg string) error {
	if _, ok := err.(gophercloud.ErrDefault404); ok {
		d.SetId("")
		return nil
	}

	return fmt.Errorf("%s %s: %s", msg, d.Id(), err)
}

// getEnvVarWithDefault retrieves the value of the environment variable
// named by the key. If the variable is not present, return the default
//value instead.
func getenvWithDefault(key string, defaultValue string) string {
	if val, exists := os.LookupEnv(key); exists {
		return val
	}
	return defaultValue
}

// set the lock state of a resource based upon a boolean
func setLockState(client *brightbox.Client, isLocked bool, resource interface{}) error {
	if isLocked {
		return client.LockResource(resource)
	}
	return client.UnLockResource(resource)
}

// strSliceContains checks if a given string is contained in a slice
// When anybody asks why Go needs generics, here you go.
func strSliceContains(haystack []string, needle string) bool {
	for _, s := range haystack {
		if s == needle {
			return true
		}
	}
	return false
}

// Check a JSON object is correct
func validateJSONObject(v interface{}, k string) ([]string, []error) {
	if v == nil || v.(string) == "" {
		return nil, []error{fmt.Errorf("%q value must not be empty", k)}
	}

	var j map[string]interface{}
	s := v.(string)

	err := json.Unmarshal([]byte(s), &j)
	if err != nil {
		return nil, []error{fmt.Errorf("%q must be a JSON object: %s", k, err)}
	}

	return nil, nil
}

func diffSuppressJSONObject(k, old, new string, d *schema.ResourceData) bool {
	if strSliceContains([]string{"{}", ""}, old) &&
		strSliceContains([]string{"{}", ""}, new) {
		return true
	}
	return false
}

// HashcodeString hashes a string to a unique hashcode.
//
// crc32 returns a uint32, but for our use we need
// and non negative integer. Here we cast to an integer
// and invert it if the result is negative.
func HashcodeString(s string) int {
	v := int(crc32.ChecksumIEEE([]byte(s)))
	if v >= 0 {
		return v
	}
	if -v >= 0 {
		return -v
	}
	// v == MinInt
	return 0
}

// StringIsValidFirewallTarget checks whether a string would
// pass the Iptables validation as a valid source or destination.
func stringIsValidFirewallTarget() schema.SchemaValidateFunc {
	return validation.Any(
		validation.StringInSlice([]string{"any"}, false),
		validation.StringMatch(serverRegexp, "must be a valid server ID"),
		validation.StringMatch(serverGroupRegexp, "must be a valid server group ID"),
		validation.StringMatch(loadBalancerRegexp, "must be a valid load balancer ID"),
		validation.IsCIDR,
		validation.IsIPAddress,
	)
}

// Get a list of server IDs from a list of servers
func serverIDListFromNodes(
	nodes []brightbox.Server,
) []string {
	nodeIds := make([]string, len(nodes))
	for i, node := range nodes {
		nodeIds[i] = node.Id
	}
	return nodeIds
}
<|endoftext|>


<commit_before>package path

import (
	"fmt"
	"net/url"
	"os"
	"path/filepath"
	"strings"

	log "github.com/sirupsen/logrus"
)

// ResolvedFilePath represents a resolved file path and intended to prevent unintentional use of not verified file path.
// It is always either a URL or an absolute path.
type ResolvedFilePath string

// resolveSymbolicLinkRecursive resolves the symlink path recursively to its
// canonical path on the file system, with a maximum nesting level of maxDepth.
// If path is not a symlink, returns the verbatim copy of path and err of nil.
func resolveSymbolicLinkRecursive(path string, maxDepth int) (string, error) {
	resolved, err := os.Readlink(path)
	if err != nil {
		// path is not a symbolic link
		_, ok := err.(*os.PathError)
		if ok {
			return path, nil
		}
		// Other error has occured
		return "", err
	}

	if maxDepth == 0 {
		return "", fmt.Errorf("maximum nesting level reached")
	}

	// If we resolved to a relative symlink, make sure we use the absolute
	// path for further resolving
	if !strings.HasPrefix(resolved, "/") {
		basePath := filepath.Dir(path)
		resolved = filepath.Join(basePath, resolved)
	}

	return resolveSymbolicLinkRecursive(resolved, maxDepth-1)
}

// isURLSchemeAllowed returns true if the protocol scheme is in the list of
// allowed URL schemes.
func isURLSchemeAllowed(scheme string, allowed []string) bool {
	isAllowed := false
	if len(allowed) > 0 {
		for _, s := range allowed {
			if strings.EqualFold(scheme, s) {
				isAllowed = true
				break
			}
		}
	}

	// Empty scheme means local file
	return isAllowed && scheme != ""
}

// ResolveFilePath will inspect and resolve given file, and make sure that its final path is within the boundaries of
// the path specified in repoRoot.
//
// appPath is the path we're operating in, e.g. where a Helm chart was unpacked
// to. repoRoot is the path to the root of the repository.
//
// If either appPath or repoRoot is relative, it will be treated as relative
// to the current working directory.
//
// valueFile is the path to a value file, relative to appPath. If valueFile is
// specified as an absolute path (i.e. leading slash), it will be treated as
// relative to the repoRoot. In case valueFile is a symlink in the extracted
// chart, it will be resolved recursively and the decision of whether it is in
// the boundary of repoRoot will be made using the final resolved path.
// valueFile can also be a remote URL with a protocol scheme as prefix,
// in which case the scheme must be included in the list of allowed schemes
// specified by allowedURLSchemes.
//
// Will return an error if either valueFile is outside the boundaries of the
// repoRoot, valueFile is an URL with a forbidden protocol scheme or if
// valueFile is a recursive symlink nested too deep. May return errors for
// other reasons as well.
//
// resolvedPath will hold the absolute, resolved path for valueFile on success
// or set to the empty string on failure.
//
// isRemote will be set to true if valueFile is an URL using an allowed
// protocol scheme, or to false if it resolved to a local file.
func ResolveFilePath(appPath, repoRoot, valueFile string, allowedURLSchemes []string) (resolvedPath ResolvedFilePath, isRemote bool, err error) {
	// We do not provide the path in the error message, because it will be
	// returned to the user and could be used for information gathering.
	// Instead, we log the concrete error details.
	resolveFailure := func(path string, err error) error {
		log.Errorf("failed to resolve path '%s': %v", path, err)
		return fmt.Errorf("internal error: failed to resolve path. Check logs for more details")
	}

	// A value file can be specified as an URL to a remote resource.
	// We only allow certain URL schemes for remote value files.
	url, err := url.Parse(valueFile)
	if err == nil {
		// If scheme is empty, it means we parsed a path only
		if url.Scheme != "" {
			if isURLSchemeAllowed(url.Scheme, allowedURLSchemes) {
				return ResolvedFilePath(valueFile), true, nil
			} else {
				return "", false, fmt.Errorf("the URL scheme '%s' is not allowed", url.Scheme)
			}
		}
	}

	// Ensure that our repository root is absolute
	absRepoPath, err := filepath.Abs(repoRoot)
	if err != nil {
		return "", false, resolveFailure(repoRoot, err)
	}

	// If the path to the file is relative, join it with the current working directory (appPath)
	// Otherwise, join it with the repository's root
	path := valueFile
	if !filepath.IsAbs(path) {
		absWorkDir, err := filepath.Abs(appPath)
		if err != nil {
			return "", false, resolveFailure(repoRoot, err)
		}
		path = filepath.Join(absWorkDir, path)
	} else {
		path = filepath.Join(absRepoPath, path)
	}

	// Ensure any symbolic link is resolved before we
	delinkedPath, err := resolveSymbolicLinkRecursive(path, 10)
	if err != nil {
		return "", false, resolveFailure(path, err)
	}
	path = delinkedPath

	// Resolve the joined path to an absolute path
	path, err = filepath.Abs(path)
	if err != nil {
		return "", false, resolveFailure(path, err)
	}

	// Ensure our root path has a trailing slash, otherwise the following check
	// would return true if root is /foo and path would be /foo2
	requiredRootPath := absRepoPath
	if !strings.HasSuffix(requiredRootPath, "/") {
		requiredRootPath += "/"
	}

	// Make sure that the resolved path to values file is within the repository's root path
	if !strings.HasPrefix(path, requiredRootPath) {
		return "", false, fmt.Errorf("value file '%s' resolved to outside repository root", valueFile)
	}

	return ResolvedFilePath(path), false, nil
}
<commit_msg>fix: Use os.PathSeparator instead of hard-coded string to resolve local file paths (#10945) (#10946)<commit_after>package path

import (
	"fmt"
	"net/url"
	"os"
	"path/filepath"
	"strings"

	log "github.com/sirupsen/logrus"
)

// ResolvedFilePath represents a resolved file path and intended to prevent unintentional use of not verified file path.
// It is always either a URL or an absolute path.
type ResolvedFilePath string

// resolveSymbolicLinkRecursive resolves the symlink path recursively to its
// canonical path on the file system, with a maximum nesting level of maxDepth.
// If path is not a symlink, returns the verbatim copy of path and err of nil.
func resolveSymbolicLinkRecursive(path string, maxDepth int) (string, error) {
	resolved, err := os.Readlink(path)
	if err != nil {
		// path is not a symbolic link
		_, ok := err.(*os.PathError)
		if ok {
			return path, nil
		}
		// Other error has occured
		return "", err
	}

	if maxDepth == 0 {
		return "", fmt.Errorf("maximum nesting level reached")
	}

	// If we resolved to a relative symlink, make sure we use the absolute
	// path for further resolving
	if !strings.HasPrefix(resolved, string(os.PathSeparator)) {
		basePath := filepath.Dir(path)
		resolved = filepath.Join(basePath, resolved)
	}

	return resolveSymbolicLinkRecursive(resolved, maxDepth-1)
}

// isURLSchemeAllowed returns true if the protocol scheme is in the list of
// allowed URL schemes.
func isURLSchemeAllowed(scheme string, allowed []string) bool {
	isAllowed := false
	if len(allowed) > 0 {
		for _, s := range allowed {
			if strings.EqualFold(scheme, s) {
				isAllowed = true
				break
			}
		}
	}

	// Empty scheme means local file
	return isAllowed && scheme != ""
}

// ResolveFilePath will inspect and resolve given file, and make sure that its final path is within the boundaries of
// the path specified in repoRoot.
//
// appPath is the path we're operating in, e.g. where a Helm chart was unpacked
// to. repoRoot is the path to the root of the repository.
//
// If either appPath or repoRoot is relative, it will be treated as relative
// to the current working directory.
//
// valueFile is the path to a value file, relative to appPath. If valueFile is
// specified as an absolute path (i.e. leading slash), it will be treated as
// relative to the repoRoot. In case valueFile is a symlink in the extracted
// chart, it will be resolved recursively and the decision of whether it is in
// the boundary of repoRoot will be made using the final resolved path.
// valueFile can also be a remote URL with a protocol scheme as prefix,
// in which case the scheme must be included in the list of allowed schemes
// specified by allowedURLSchemes.
//
// Will return an error if either valueFile is outside the boundaries of the
// repoRoot, valueFile is an URL with a forbidden protocol scheme or if
// valueFile is a recursive symlink nested too deep. May return errors for
// other reasons as well.
//
// resolvedPath will hold the absolute, resolved path for valueFile on success
// or set to the empty string on failure.
//
// isRemote will be set to true if valueFile is an URL using an allowed
// protocol scheme, or to false if it resolved to a local file.
func ResolveFilePath(appPath, repoRoot, valueFile string, allowedURLSchemes []string) (resolvedPath ResolvedFilePath, isRemote bool, err error) {
	// We do not provide the path in the error message, because it will be
	// returned to the user and could be used for information gathering.
	// Instead, we log the concrete error details.
	resolveFailure := func(path string, err error) error {
		log.Errorf("failed to resolve path '%s': %v", path, err)
		return fmt.Errorf("internal error: failed to resolve path. Check logs for more details")
	}

	// A value file can be specified as an URL to a remote resource.
	// We only allow certain URL schemes for remote value files.
	url, err := url.Parse(valueFile)
	if err == nil {
		// If scheme is empty, it means we parsed a path only
		if url.Scheme != "" {
			if isURLSchemeAllowed(url.Scheme, allowedURLSchemes) {
				return ResolvedFilePath(valueFile), true, nil
			} else {
				return "", false, fmt.Errorf("the URL scheme '%s' is not allowed", url.Scheme)
			}
		}
	}

	// Ensure that our repository root is absolute
	absRepoPath, err := filepath.Abs(repoRoot)
	if err != nil {
		return "", false, resolveFailure(repoRoot, err)
	}

	// If the path to the file is relative, join it with the current working directory (appPath)
	// Otherwise, join it with the repository's root
	path := valueFile
	if !filepath.IsAbs(path) {
		absWorkDir, err := filepath.Abs(appPath)
		if err != nil {
			return "", false, resolveFailure(repoRoot, err)
		}
		path = filepath.Join(absWorkDir, path)
	} else {
		path = filepath.Join(absRepoPath, path)
	}

	// Ensure any symbolic link is resolved before we
	delinkedPath, err := resolveSymbolicLinkRecursive(path, 10)
	if err != nil {
		return "", false, resolveFailure(path, err)
	}
	path = delinkedPath

	// Resolve the joined path to an absolute path
	path, err = filepath.Abs(path)
	if err != nil {
		return "", false, resolveFailure(path, err)
	}

	// Ensure our root path has a trailing slash, otherwise the following check
	// would return true if root is /foo and path would be /foo2
	requiredRootPath := absRepoPath
	if !strings.HasSuffix(requiredRootPath, string(os.PathSeparator)) {
		requiredRootPath += string(os.PathSeparator)
	}

	// Make sure that the resolved path to values file is within the repository's root path
	if !strings.HasPrefix(path, requiredRootPath) {
		return "", false, fmt.Errorf("value file '%s' resolved to outside repository root", valueFile)
	}

	return ResolvedFilePath(path), false, nil
}
<|endoftext|>


<commit_before>package libkbfs

import (
	"fmt"
	"sync"
	"time"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/protocol/go"
	"golang.org/x/net/context"
)

// MDServerRemote is an implementation of the MDServer interface.
type MDServerRemote struct {
	config Config
	conn   *Connection

	observerMu sync.Mutex // protects observers
	observers  map[TlfID]chan<- error

	testClient keybase1.GenericClient // for testing
}

// Test that MDServerRemote fully implements the MDServer interface.
var _ MDServer = (*MDServerRemote)(nil)

// Test that MDServerRemote fully implements the KeyServer interface.
var _ KeyServer = (*MDServerRemote)(nil)

// NewMDServerRemote returns a new instance of MDServerRemote.
func NewMDServerRemote(ctx context.Context, config Config, srvAddr string) *MDServerRemote {
	mdServer := &MDServerRemote{
		config:    config,
		observers: make(map[TlfID]chan<- error),
	}
	connection := NewConnection(ctx, config, srvAddr, mdServer, MDServerUnwrapError)
	mdServer.conn = connection
	return mdServer
}

// For testing.
func newMDServerRemoteWithClient(ctx context.Context, config Config,
	testClient keybase1.GenericClient) *MDServerRemote {
	mdServer := &MDServerRemote{config: config, testClient: testClient}
	return mdServer
}

// OnConnect implements the ConnectionHandler interface.
func (md *MDServerRemote) OnConnect(ctx context.Context,
	conn *Connection, client keybase1.GenericClient) error {
	// get UID, deviceKID and session token
	var err error
	var user keybase1.UID
	user, err = md.config.KBPKI().GetLoggedInUser(ctx)
	if err != nil {
		return err
	}
	var key CryptPublicKey
	key, err = md.config.KBPKI().GetCurrentCryptPublicKey(ctx)
	if err != nil {
		return err
	}
	var token string
	var session *libkb.Session
	session, err = md.config.KBPKI().GetSession(ctx)
	if err != nil {
		libkb.G.Log.Warning("MDServerRemote: error getting session %q", err)
		return err
	} else if session != nil {
		token = session.GetToken()
	}

	// authenticate
	creds := keybase1.AuthenticateArg{
		User:      user,
		DeviceKID: key.KID,
		Sid:       token,
	}

	// save the conn pointer
	md.conn = conn

	// using conn.DoCommand here would cause problematic recursion
	return runUnlessCanceled(ctx, func() error {
		c := keybase1.MetadataClient{Cli: client}
		return c.Authenticate(creds)
	})
}

// OnConnectError implements the ConnectionHandler interface.
func (md *MDServerRemote) OnConnectError(err error, wait time.Duration) {
	libkb.G.Log.Warning("MDServerRemote: connection error: %q; retrying in %s",
		err, wait)
	// TODO: it might make sense to show something to the user if this is
	// due to authentication, for example.
	md.cancelObservers()
}

// OnDisconnected implements the ConnectionHandler interface.
func (md *MDServerRemote) OnDisconnected() {
	md.cancelObservers()
}

// Signal errors and clear any registered observers.
func (md *MDServerRemote) cancelObservers() {
	md.observerMu.Lock()
	defer md.observerMu.Unlock()
	// fire errors for any registered observers
	for id, observerChan := range md.observers {
		md.signalObserverLocked(observerChan, id, MDServerDisconnected{})
	}
}

// Signal an observer. The observer lock must be held.
func (md *MDServerRemote) signalObserverLocked(observerChan chan<- error, id TlfID, err error) {
	observerChan <- err
	close(observerChan)
	delete(md.observers, id)
}

// Helper to return a metadata client.
func (md *MDServerRemote) client() keybase1.MetadataClient {
	if md.testClient != nil {
		// for testing
		return keybase1.MetadataClient{Cli: md.testClient}
	}
	return keybase1.MetadataClient{Cli: md.conn.GetClient()}
}

// Helper to call an rpc command.
func (md *MDServerRemote) doCommand(ctx context.Context, command func() error) error {
	if md.testClient != nil {
		// for testing
		return runUnlessCanceled(ctx, command)
	}
	return md.conn.DoCommand(ctx, command)
}

// Helper used to retrieve metadata blocks from the MD server.
func (md *MDServerRemote) get(ctx context.Context, id TlfID, handle *TlfHandle,
	unmerged bool, start, stop MetadataRevision) (TlfID, []*RootMetadataSigned, error) {
	// figure out which args to send
	if id == NullTlfID && handle == nil {
		return id, nil, MDInvalidGetArguments{
			id:     id,
			handle: handle,
		}
	}
	arg := keybase1.GetMetadataArg{
		StartRevision: start.Number(),
		StopRevision:  stop.Number(),
		Unmerged:      unmerged,
	}
	if id == NullTlfID {
		arg.FolderHandle = handle.ToBytes(md.config)
	} else {
		arg.FolderID = id.String()
	}

	// request
	var err error
	var response keybase1.MetadataResponse
	err = md.doCommand(ctx, func() error {
		response, err = md.client().GetMetadata(arg)
		return err
	})
	if err != nil {
		return id, nil, err
	}

	// response
	id = ParseTlfID(response.FolderID)
	if id == NullTlfID {
		return id, nil, MDInvalidTlfID{response.FolderID}
	}

	// deserialize blocks
	rmdses := make([]*RootMetadataSigned, len(response.MdBlocks))
	for i := range response.MdBlocks {
		var rmds RootMetadataSigned
		err = md.config.Codec().Decode(response.MdBlocks[i], &rmds)
		if err != nil {
			return id, rmdses, err
		}
		rmdses[i] = &rmds
	}
	return id, rmdses, nil
}

// GetForHandle implements the MDServer interface for MDServerRemote.
func (md *MDServerRemote) GetForHandle(ctx context.Context, handle *TlfHandle, unmerged bool) (
	TlfID, *RootMetadataSigned, error) {
	id, rmdses, err := md.get(ctx, NullTlfID, handle, unmerged,
		MetadataRevisionUninitialized, MetadataRevisionUninitialized)
	if err != nil {
		return id, nil, err
	}
	if len(rmdses) == 0 {
		return id, nil, nil
	}
	return id, rmdses[0], nil
}

// GetForTLF implements the MDServer interface for MDServerRemote.
func (md *MDServerRemote) GetForTLF(ctx context.Context, id TlfID, unmerged bool) (
	*RootMetadataSigned, error) {
	_, rmdses, err := md.get(ctx, id, nil, unmerged,
		MetadataRevisionUninitialized, MetadataRevisionUninitialized)
	if err != nil {
		return nil, err
	}
	if len(rmdses) == 0 {
		return nil, nil
	}
	return rmdses[0], nil
}

// GetRange implements the MDServer interface for MDServerRemote.
func (md *MDServerRemote) GetRange(ctx context.Context, id TlfID, unmerged bool,
	start, stop MetadataRevision) ([]*RootMetadataSigned, error) {
	_, rmds, err := md.get(ctx, id, nil, unmerged, start, stop)
	return rmds, err
}

// Put implements the MDServer interface for MDServerRemote.
func (md *MDServerRemote) Put(ctx context.Context, rmds *RootMetadataSigned) error {
	// encode MD block
	rmdsBytes, err := md.config.Codec().Encode(rmds)
	if err != nil {
		return err
	}
	// put request
	return md.doCommand(ctx, func() error {
		return md.client().PutMetadata(rmdsBytes)
	})
}

// PruneUnmerged implements the MDServer interface for MDServerRemote.
func (md *MDServerRemote) PruneUnmerged(ctx context.Context, id TlfID) error {
	return md.doCommand(ctx, func() error {
		return md.client().PruneUnmerged(id.String())
	})
}

// MetadataUpdate implements the MetadataUpdateProtocol interface.
func (md *MDServerRemote) MetadataUpdate(arg keybase1.MetadataUpdateArg) error {
	id := ParseTlfID(arg.FolderID)
	if id == NullTlfID {
		return MDServerErrorBadRequest{"Invalid folder ID"}
	}

	md.observerMu.Lock()
	defer md.observerMu.Unlock()
	observerChan, ok := md.observers[id]
	if !ok {
		// not registered
		return nil
	}

	// signal that we've seen the update
	md.signalObserverLocked(observerChan, id, nil)
	return nil
}

// RegisterForUpdate implements the MDServer interface for MDServerRemote.
func (md *MDServerRemote) RegisterForUpdate(ctx context.Context, id TlfID,
	currHead MetadataRevision) (<-chan error, error) {
	arg := keybase1.RegisterForUpdatesArg{
		FolderID:     id.String(),
		CurrRevision: currHead.Number(),
	}

	// setup the server to receive updates
	err := md.conn.Serve(keybase1.MetadataUpdateProtocol(md))
	if err != nil {
		libkb.G.Log.Warning("MDServerRemote: unable to create update server %q", err)
		return nil, err
	}

	// register
	var c chan error
	err = md.doCommand(ctx, func() error {
		// keep re-adding the observer on retries.
		// OnDisconnected/OnConnectError will remove it.
		func() {
			md.observerMu.Lock()
			defer md.observerMu.Unlock()
			if _, ok := md.observers[id]; ok {
				panic(fmt.Sprintf("Attempted double-registration for MDServerRemote %p",
					md))
			}
			c = make(chan error, 1)
			md.observers[id] = c
		}()
		return md.client().RegisterForUpdates(arg)
	})
	if err != nil {
		// cancel the observer if it still exists
		md.observerMu.Lock()
		defer md.observerMu.Unlock()
		if observerChan, ok := md.observers[id]; ok {
			delete(md.observers, id)
			close(observerChan)
		}
		c = nil
	}

	return c, err
}

// Shutdown implements the MDServer interface for MDServerRemote.
func (md *MDServerRemote) Shutdown() {
	// close the connection
	md.conn.Shutdown()
	// cancel pending observers
	md.cancelObservers()
}

//
// The below methods support the MD server acting as the key server.
// This will be the case for v1 of KBFS but we may move to our own
// separate key server at some point.
//

// GetTLFCryptKeyServerHalf is an implementation of the KeyServer interface.
func (md *MDServerRemote) GetTLFCryptKeyServerHalf(ctx context.Context,
	serverHalfID TLFCryptKeyServerHalfID) (TLFCryptKeyServerHalf, error) {
	// encode the ID
	idBytes, err := md.config.Codec().Encode(serverHalfID)
	if err != nil {
		return TLFCryptKeyServerHalf{}, err
	}

	// get the key
	var keyBytes []byte
	err = md.doCommand(ctx, func() error {
		keyBytes, err = md.client().GetKey(idBytes)
		return err
	})
	if err != nil {
		return TLFCryptKeyServerHalf{}, err
	}

	// decode the key
	var serverHalf TLFCryptKeyServerHalf
	err = md.config.Codec().Decode(keyBytes, &serverHalf)
	if err != nil {
		return TLFCryptKeyServerHalf{}, err
	}

	return serverHalf, nil
}

// PutTLFCryptKeyServerHalves is an implementation of the KeyServer interface.
func (md *MDServerRemote) PutTLFCryptKeyServerHalves(ctx context.Context,
	serverKeyHalves map[keybase1.UID]map[keybase1.KID]TLFCryptKeyServerHalf) error {
	// flatten out the map into an array
	var keyHalves []keybase1.KeyHalf
	for user, deviceMap := range serverKeyHalves {
		for deviceKID, serverHalf := range deviceMap {
			keyHalf, err := md.config.Codec().Encode(serverHalf)
			if err != nil {
				return err
			}
			keyHalves = append(keyHalves,
				keybase1.KeyHalf{
					User:      user,
					DeviceKID: deviceKID,
					Key:       keyHalf,
				})
		}
	}
	// put the keys
	return md.doCommand(ctx, func() error {
		return md.client().PutKeys(keyHalves)
	})
}
<commit_msg>mdserver_remote: use logger.Logger<commit_after>package libkbfs

import (
	"fmt"
	"sync"
	"time"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	keybase1 "github.com/keybase/client/protocol/go"
	"golang.org/x/net/context"
)

// MDServerRemote is an implementation of the MDServer interface.
type MDServerRemote struct {
	config Config
	conn   *Connection
	log    logger.Logger

	observerMu sync.Mutex // protects observers
	observers  map[TlfID]chan<- error

	testClient keybase1.GenericClient // for testing
}

// Test that MDServerRemote fully implements the MDServer interface.
var _ MDServer = (*MDServerRemote)(nil)

// Test that MDServerRemote fully implements the KeyServer interface.
var _ KeyServer = (*MDServerRemote)(nil)

// NewMDServerRemote returns a new instance of MDServerRemote.
func NewMDServerRemote(ctx context.Context, config Config, srvAddr string) *MDServerRemote {
	mdServer := &MDServerRemote{
		config:    config,
		observers: make(map[TlfID]chan<- error),
		log:       config.MakeLogger(""),
	}
	connection := NewConnection(ctx, config, srvAddr, mdServer, MDServerUnwrapError)
	mdServer.conn = connection
	return mdServer
}

// For testing.
func newMDServerRemoteWithClient(ctx context.Context, config Config,
	testClient keybase1.GenericClient) *MDServerRemote {
	mdServer := &MDServerRemote{
		config:     config,
		testClient: testClient,
		log:        config.MakeLogger(""),
	}
	return mdServer
}

// OnConnect implements the ConnectionHandler interface.
func (md *MDServerRemote) OnConnect(ctx context.Context,
	conn *Connection, client keybase1.GenericClient) error {
	// get UID, deviceKID and session token
	var err error
	var user keybase1.UID
	user, err = md.config.KBPKI().GetLoggedInUser(ctx)
	if err != nil {
		return err
	}
	var key CryptPublicKey
	key, err = md.config.KBPKI().GetCurrentCryptPublicKey(ctx)
	if err != nil {
		return err
	}
	var token string
	var session *libkb.Session
	session, err = md.config.KBPKI().GetSession(ctx)
	if err != nil {
		md.log.CWarningf(ctx, "MDServerRemote: error getting session %q", err)
		return err
	} else if session != nil {
		token = session.GetToken()
	}

	// authenticate
	creds := keybase1.AuthenticateArg{
		User:      user,
		DeviceKID: key.KID,
		Sid:       token,
	}

	// save the conn pointer
	md.conn = conn

	// using conn.DoCommand here would cause problematic recursion
	return runUnlessCanceled(ctx, func() error {
		c := keybase1.MetadataClient{Cli: client}
		return c.Authenticate(creds)
	})
}

// OnConnectError implements the ConnectionHandler interface.
func (md *MDServerRemote) OnConnectError(err error, wait time.Duration) {
	md.log.Warning("MDServerRemote: connection error: %q; retrying in %s",
		err, wait)
	// TODO: it might make sense to show something to the user if this is
	// due to authentication, for example.
	md.cancelObservers()
}

// OnDisconnected implements the ConnectionHandler interface.
func (md *MDServerRemote) OnDisconnected() {
	md.cancelObservers()
}

// Signal errors and clear any registered observers.
func (md *MDServerRemote) cancelObservers() {
	md.observerMu.Lock()
	defer md.observerMu.Unlock()
	// fire errors for any registered observers
	for id, observerChan := range md.observers {
		md.signalObserverLocked(observerChan, id, MDServerDisconnected{})
	}
}

// Signal an observer. The observer lock must be held.
func (md *MDServerRemote) signalObserverLocked(observerChan chan<- error, id TlfID, err error) {
	observerChan <- err
	close(observerChan)
	delete(md.observers, id)
}

// Helper to return a metadata client.
func (md *MDServerRemote) client() keybase1.MetadataClient {
	if md.testClient != nil {
		// for testing
		return keybase1.MetadataClient{Cli: md.testClient}
	}
	return keybase1.MetadataClient{Cli: md.conn.GetClient()}
}

// Helper to call an rpc command.
func (md *MDServerRemote) doCommand(ctx context.Context, command func() error) error {
	if md.testClient != nil {
		// for testing
		return runUnlessCanceled(ctx, command)
	}
	return md.conn.DoCommand(ctx, command)
}

// Helper used to retrieve metadata blocks from the MD server.
func (md *MDServerRemote) get(ctx context.Context, id TlfID, handle *TlfHandle,
	unmerged bool, start, stop MetadataRevision) (TlfID, []*RootMetadataSigned, error) {
	// figure out which args to send
	if id == NullTlfID && handle == nil {
		return id, nil, MDInvalidGetArguments{
			id:     id,
			handle: handle,
		}
	}
	arg := keybase1.GetMetadataArg{
		StartRevision: start.Number(),
		StopRevision:  stop.Number(),
		Unmerged:      unmerged,
	}
	if id == NullTlfID {
		arg.FolderHandle = handle.ToBytes(md.config)
	} else {
		arg.FolderID = id.String()
	}

	// request
	var err error
	var response keybase1.MetadataResponse
	err = md.doCommand(ctx, func() error {
		response, err = md.client().GetMetadata(arg)
		return err
	})
	if err != nil {
		return id, nil, err
	}

	// response
	id = ParseTlfID(response.FolderID)
	if id == NullTlfID {
		return id, nil, MDInvalidTlfID{response.FolderID}
	}

	// deserialize blocks
	rmdses := make([]*RootMetadataSigned, len(response.MdBlocks))
	for i := range response.MdBlocks {
		var rmds RootMetadataSigned
		err = md.config.Codec().Decode(response.MdBlocks[i], &rmds)
		if err != nil {
			return id, rmdses, err
		}
		rmdses[i] = &rmds
	}
	return id, rmdses, nil
}

// GetForHandle implements the MDServer interface for MDServerRemote.
func (md *MDServerRemote) GetForHandle(ctx context.Context, handle *TlfHandle, unmerged bool) (
	TlfID, *RootMetadataSigned, error) {
	id, rmdses, err := md.get(ctx, NullTlfID, handle, unmerged,
		MetadataRevisionUninitialized, MetadataRevisionUninitialized)
	if err != nil {
		return id, nil, err
	}
	if len(rmdses) == 0 {
		return id, nil, nil
	}
	return id, rmdses[0], nil
}

// GetForTLF implements the MDServer interface for MDServerRemote.
func (md *MDServerRemote) GetForTLF(ctx context.Context, id TlfID, unmerged bool) (
	*RootMetadataSigned, error) {
	_, rmdses, err := md.get(ctx, id, nil, unmerged,
		MetadataRevisionUninitialized, MetadataRevisionUninitialized)
	if err != nil {
		return nil, err
	}
	if len(rmdses) == 0 {
		return nil, nil
	}
	return rmdses[0], nil
}

// GetRange implements the MDServer interface for MDServerRemote.
func (md *MDServerRemote) GetRange(ctx context.Context, id TlfID, unmerged bool,
	start, stop MetadataRevision) ([]*RootMetadataSigned, error) {
	_, rmds, err := md.get(ctx, id, nil, unmerged, start, stop)
	return rmds, err
}

// Put implements the MDServer interface for MDServerRemote.
func (md *MDServerRemote) Put(ctx context.Context, rmds *RootMetadataSigned) error {
	// encode MD block
	rmdsBytes, err := md.config.Codec().Encode(rmds)
	if err != nil {
		return err
	}
	// put request
	return md.doCommand(ctx, func() error {
		return md.client().PutMetadata(rmdsBytes)
	})
}

// PruneUnmerged implements the MDServer interface for MDServerRemote.
func (md *MDServerRemote) PruneUnmerged(ctx context.Context, id TlfID) error {
	return md.doCommand(ctx, func() error {
		return md.client().PruneUnmerged(id.String())
	})
}

// MetadataUpdate implements the MetadataUpdateProtocol interface.
func (md *MDServerRemote) MetadataUpdate(arg keybase1.MetadataUpdateArg) error {
	id := ParseTlfID(arg.FolderID)
	if id == NullTlfID {
		return MDServerErrorBadRequest{"Invalid folder ID"}
	}

	md.observerMu.Lock()
	defer md.observerMu.Unlock()
	observerChan, ok := md.observers[id]
	if !ok {
		// not registered
		return nil
	}

	// signal that we've seen the update
	md.signalObserverLocked(observerChan, id, nil)
	return nil
}

// RegisterForUpdate implements the MDServer interface for MDServerRemote.
func (md *MDServerRemote) RegisterForUpdate(ctx context.Context, id TlfID,
	currHead MetadataRevision) (<-chan error, error) {
	arg := keybase1.RegisterForUpdatesArg{
		FolderID:     id.String(),
		CurrRevision: currHead.Number(),
	}

	// setup the server to receive updates
	err := md.conn.Serve(keybase1.MetadataUpdateProtocol(md))
	if err != nil {
		md.log.CWarningf(ctx,
			"MDServerRemote: unable to create update server %q", err)
		return nil, err
	}

	// register
	var c chan error
	err = md.doCommand(ctx, func() error {
		// keep re-adding the observer on retries.
		// OnDisconnected/OnConnectError will remove it.
		func() {
			md.observerMu.Lock()
			defer md.observerMu.Unlock()
			if _, ok := md.observers[id]; ok {
				panic(fmt.Sprintf("Attempted double-registration for MDServerRemote %p",
					md))
			}
			c = make(chan error, 1)
			md.observers[id] = c
		}()
		return md.client().RegisterForUpdates(arg)
	})
	if err != nil {
		// cancel the observer if it still exists
		md.observerMu.Lock()
		defer md.observerMu.Unlock()
		if observerChan, ok := md.observers[id]; ok {
			delete(md.observers, id)
			close(observerChan)
		}
		c = nil
	}

	return c, err
}

// Shutdown implements the MDServer interface for MDServerRemote.
func (md *MDServerRemote) Shutdown() {
	// close the connection
	md.conn.Shutdown()
	// cancel pending observers
	md.cancelObservers()
}

//
// The below methods support the MD server acting as the key server.
// This will be the case for v1 of KBFS but we may move to our own
// separate key server at some point.
//

// GetTLFCryptKeyServerHalf is an implementation of the KeyServer interface.
func (md *MDServerRemote) GetTLFCryptKeyServerHalf(ctx context.Context,
	serverHalfID TLFCryptKeyServerHalfID) (TLFCryptKeyServerHalf, error) {
	// encode the ID
	idBytes, err := md.config.Codec().Encode(serverHalfID)
	if err != nil {
		return TLFCryptKeyServerHalf{}, err
	}

	// get the key
	var keyBytes []byte
	err = md.doCommand(ctx, func() error {
		keyBytes, err = md.client().GetKey(idBytes)
		return err
	})
	if err != nil {
		return TLFCryptKeyServerHalf{}, err
	}

	// decode the key
	var serverHalf TLFCryptKeyServerHalf
	err = md.config.Codec().Decode(keyBytes, &serverHalf)
	if err != nil {
		return TLFCryptKeyServerHalf{}, err
	}

	return serverHalf, nil
}

// PutTLFCryptKeyServerHalves is an implementation of the KeyServer interface.
func (md *MDServerRemote) PutTLFCryptKeyServerHalves(ctx context.Context,
	serverKeyHalves map[keybase1.UID]map[keybase1.KID]TLFCryptKeyServerHalf) error {
	// flatten out the map into an array
	var keyHalves []keybase1.KeyHalf
	for user, deviceMap := range serverKeyHalves {
		for deviceKID, serverHalf := range deviceMap {
			keyHalf, err := md.config.Codec().Encode(serverHalf)
			if err != nil {
				return err
			}
			keyHalves = append(keyHalves,
				keybase1.KeyHalf{
					User:      user,
					DeviceKID: deviceKID,
					Key:       keyHalf,
				})
		}
	}
	// put the keys
	return md.doCommand(ctx, func() error {
		return md.client().PutKeys(keyHalves)
	})
}
<|endoftext|>


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

// +build acceptance

package acceptance

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"testing"
	"time"

	"github.com/cockroachdb/cockroach/acceptance/localcluster"
	"github.com/cockroachdb/cockroach/proto"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/retry"
)

type details struct {
	NodeID proto.NodeID `json:"nodeID"`
}

var retryOptions = retry.Options{
	InitialBackoff: 100 * time.Millisecond,
	MaxRetries:     4,
	Multiplier:     2,
}

// get performs an HTTPS GET to the specified path for a specific node.
func get(t *testing.T, client *http.Client, node *localcluster.Container, path string) []byte {
	url := fmt.Sprintf("https://%s%s", node.Addr(""), path)
	// TODO(bram) #2059: Remove retry logic.
	for r := retry.Start(retryOptions); r.Next(); {
		resp, err := client.Get(url)
		if err != nil {
			t.Logf("could not GET %s - %s", url, err)
			continue
		}
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			t.Logf("could not read body for %s - %s", url, err)
			continue
		}
		defer resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			t.Logf("could not GET %s - statuscode: %d - body: %s", url, resp.StatusCode, body)
			continue
		}
		t.Logf("OK response from %s", url)
		return body
	}
	t.Fatalf("There was an error retrieving %s", url)
	return []byte("")
}

// checkNode checks all the endpoints of the status server hosted by node and
// requests info for the node with otherNodeID. That node could be the same
// other node, the same node or "local".
func checkNode(t *testing.T, client *http.Client, node *localcluster.Container, nodeID, otherNodeID, expectedNodeID string) {
	body := get(t, client, node, "/_status/details/"+otherNodeID)
	var detail details
	if err := json.Unmarshal(body, &detail); err != nil {
		t.Fatal(util.ErrorfSkipFrames(1, "unable to parse details - %s", err))
	}
	if actualNodeID := detail.NodeID.String(); actualNodeID != expectedNodeID {
		t.Fatal(util.ErrorfSkipFrames(1, "%s calling %s: node ids don't match - expected %s, actual %s", nodeID, otherNodeID, expectedNodeID, actualNodeID))
	}

	get(t, client, node, fmt.Sprintf("/_status/gossip/%s", otherNodeID))
	get(t, client, node, fmt.Sprintf("/_status/logfiles/%s", otherNodeID))
	get(t, client, node, fmt.Sprintf("/_status/logs/%s", otherNodeID))
	get(t, client, node, fmt.Sprintf("/_status/stacks/%s", otherNodeID))
	get(t, client, node, fmt.Sprintf("/_status/nodes/%s", otherNodeID))
}

// TestStatusServer starts up an N node cluster and tests the status server on
// each node.
func TestStatusServer(t *testing.T) {
	l := localcluster.Create(*numNodes, stopper)
	l.ForceLogging = true
	l.Start()
	defer l.Stop()
	checkRangeReplication(t, l, 20*time.Second)

	client := &http.Client{
		Timeout: 200 * time.Millisecond,
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{
				InsecureSkipVerify: true,
			},
		},
	}

	// Get the ids for each node.
	idMap := make(map[string]string)
	for _, node := range l.Nodes {
		body := get(t, client, node, "/_status/details/local")
		var detail details
		if err := json.Unmarshal(body, &detail); err != nil {
			t.Fatalf("unable to parse details - %s", err)
		}
		idMap[node.ID] = detail.NodeID.String()
	}

	// Check local response for the every node.
	for _, node := range l.Nodes {
		checkNode(t, client, node, idMap[node.ID], "local", idMap[node.ID])
		get(t, client, node, "/_status/nodes")
		get(t, client, node, "/_status/stores")
	}

	// Proxy from the first node to the last node.
	firstNode := l.Nodes[0]
	lastNode := l.Nodes[len(l.Nodes)-1]
	firstID := idMap[firstNode.ID]
	lastID := idMap[lastNode.ID]
	checkNode(t, client, firstNode, firstID, lastID, lastID)

	// And from the last node to the first node.
	checkNode(t, client, lastNode, lastID, firstID, firstID)

	// And from the last node to the last node.
	checkNode(t, client, lastNode, lastID, lastID, lastID)
}
<commit_msg>Fix defer Close()<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.
//
// Author: Bram Gruneir (bram+code@cockroachlabs.com)

// +build acceptance

package acceptance

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"testing"
	"time"

	"github.com/cockroachdb/cockroach/acceptance/localcluster"
	"github.com/cockroachdb/cockroach/proto"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/retry"
)

type details struct {
	NodeID proto.NodeID `json:"nodeID"`
}

var retryOptions = retry.Options{
	InitialBackoff: 100 * time.Millisecond,
	MaxRetries:     4,
	Multiplier:     2,
}

// get performs an HTTPS GET to the specified path for a specific node.
func get(t *testing.T, client *http.Client, node *localcluster.Container, path string) []byte {
	url := fmt.Sprintf("https://%s%s", node.Addr(""), path)
	// TODO(bram) #2059: Remove retry logic.
	for r := retry.Start(retryOptions); r.Next(); {
		resp, err := client.Get(url)
		if err != nil {
			t.Logf("could not GET %s - %s", url, err)
			continue
		}
		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			t.Logf("could not read body for %s - %s", url, err)
			continue
		}
		if resp.StatusCode != http.StatusOK {
			t.Logf("could not GET %s - statuscode: %d - body: %s", url, resp.StatusCode, body)
			continue
		}
		t.Logf("OK response from %s", url)
		return body
	}
	t.Fatalf("There was an error retrieving %s", url)
	return []byte("")
}

// checkNode checks all the endpoints of the status server hosted by node and
// requests info for the node with otherNodeID. That node could be the same
// other node, the same node or "local".
func checkNode(t *testing.T, client *http.Client, node *localcluster.Container, nodeID, otherNodeID, expectedNodeID string) {
	body := get(t, client, node, "/_status/details/"+otherNodeID)
	var detail details
	if err := json.Unmarshal(body, &detail); err != nil {
		t.Fatal(util.ErrorfSkipFrames(1, "unable to parse details - %s", err))
	}
	if actualNodeID := detail.NodeID.String(); actualNodeID != expectedNodeID {
		t.Fatal(util.ErrorfSkipFrames(1, "%s calling %s: node ids don't match - expected %s, actual %s", nodeID, otherNodeID, expectedNodeID, actualNodeID))
	}

	get(t, client, node, fmt.Sprintf("/_status/gossip/%s", otherNodeID))
	get(t, client, node, fmt.Sprintf("/_status/logfiles/%s", otherNodeID))
	get(t, client, node, fmt.Sprintf("/_status/logs/%s", otherNodeID))
	get(t, client, node, fmt.Sprintf("/_status/stacks/%s", otherNodeID))
	get(t, client, node, fmt.Sprintf("/_status/nodes/%s", otherNodeID))
}

// TestStatusServer starts up an N node cluster and tests the status server on
// each node.
func TestStatusServer(t *testing.T) {
	l := localcluster.Create(*numNodes, stopper)
	l.ForceLogging = true
	l.Start()
	defer l.Stop()
	checkRangeReplication(t, l, 20*time.Second)

	client := &http.Client{
		Timeout: 200 * time.Millisecond,
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{
				InsecureSkipVerify: true,
			},
		},
	}

	// Get the ids for each node.
	idMap := make(map[string]string)
	for _, node := range l.Nodes {
		body := get(t, client, node, "/_status/details/local")
		var detail details
		if err := json.Unmarshal(body, &detail); err != nil {
			t.Fatalf("unable to parse details - %s", err)
		}
		idMap[node.ID] = detail.NodeID.String()
	}

	// Check local response for the every node.
	for _, node := range l.Nodes {
		checkNode(t, client, node, idMap[node.ID], "local", idMap[node.ID])
		get(t, client, node, "/_status/nodes")
		get(t, client, node, "/_status/stores")
	}

	// Proxy from the first node to the last node.
	firstNode := l.Nodes[0]
	lastNode := l.Nodes[len(l.Nodes)-1]
	firstID := idMap[firstNode.ID]
	lastID := idMap[lastNode.ID]
	checkNode(t, client, firstNode, firstID, lastID, lastID)

	// And from the last node to the first node.
	checkNode(t, client, lastNode, lastID, firstID, firstID)

	// And from the last node to the last node.
	checkNode(t, client, lastNode, lastID, lastID, lastID)
}
<|endoftext|>


<commit_before>// 结果收集与输出
package collector

import (
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/admpub/spider/app/pipeline/collector/data"
	"github.com/admpub/spider/app/spider"
	"github.com/admpub/spider/config"
	"github.com/admpub/spider/logs"
	"github.com/admpub/spider/runtime/cache"
)

// 结果收集与输出
type Collector struct {
	*spider.Spider                    //绑定的采集规则
	DataChan       chan data.DataCell //文本数据收集通道
	FileChan       chan data.FileCell //文件收集通道
	dataDocker     []data.DataCell    //分批输出结果缓存
	outType        string             //输出方式
	fileOutType    string             //文件输出方式
	FileOutPath    string             //文件输出路径
	// size     [2]uint64 //数据总输出流量统计[文本，文件]，文本暂时未统计
	dataBatch   uint64 //当前文本输出批次
	fileBatch   uint64 //当前文件输出批次
	wait        sync.WaitGroup
	sum         [4]uint64 //收集的数据总数[上次输出后文本总数，本次输出后文本总数，上次输出后文件总数，本次输出后文件总数]，非并发安全
	dataSumLock sync.RWMutex
	fileSumLock sync.RWMutex
	logger      logs.Logs
	dockerCap   int
}

func NewCollector(sp *spider.Spider) *Collector {
	var self = &Collector{}
	self.Spider = sp
	if len(sp.OutType) > 0 {
		self.outType = sp.OutType
	} else {
		self.outType = cache.Task.OutType
	}
	if len(sp.FileOutType) > 0 {
		self.fileOutType = sp.FileOutType
	} else {
		self.fileOutType = cache.Task.FileOutType
	}
	if len(self.fileOutType) == 0 {
		self.fileOutType = `local`
	}
	if len(sp.FileOutPath) > 0 {
		self.FileOutPath = sp.FileOutPath
	} else {
		self.FileOutPath = cache.Task.FileOutPath
	}
	if len(self.FileOutPath) == 0 {
		panic(`self.FileOutPath`)
		self.FileOutPath = config.FILE_DIR
	}
	if sp.DockerCap > 0 {
		self.dockerCap = sp.DockerCap
	} else {
		if cache.Task.DockerCap < 1 {
			cache.Task.DockerCap = 1
		}
		self.dockerCap = cache.Task.DockerCap
	}
	self.DataChan = make(chan data.DataCell, self.dockerCap)
	self.FileChan = make(chan data.FileCell, self.dockerCap)
	self.dataDocker = make([]data.DataCell, 0, self.dockerCap)
	self.sum = [4]uint64{}
	// self.size = [2]uint64{}
	self.dataBatch = 0
	self.fileBatch = 0
	return self
}

func (self *Collector) CollectData(dataCell data.DataCell) error {
	var err error
	defer func() {
		if recover() != nil {
			err = fmt.Errorf("输出协程已终止")
		}
	}()
	self.DataChan <- dataCell
	return err
}

func (self *Collector) CollectFile(fileCell data.FileCell) error {
	var err error
	defer func() {
		if recover() != nil {
			err = fmt.Errorf("输出协程已终止")
		}
	}()
	self.FileChan <- fileCell
	return err
}

// 停止
func (self *Collector) Stop() {
	go func() {
		defer func() {
			recover()
		}()
		close(self.DataChan)
	}()
	go func() {
		defer func() {
			recover()
		}()
		close(self.FileChan)
	}()
}

// 启动数据收集/输出管道
func (self *Collector) Start() {
	// 启动输出协程
	go self.Run()
}

// 启动数据收集/输出管道
func (self *Collector) Run() {
	dataStop := make(chan bool)
	fileStop := make(chan bool)

	go func() {
		defer func() {
			recover()
		}()
		for data := range self.DataChan {
			// 缓存分批数据
			self.dataDocker = append(self.dataDocker, data)

			// 未达到设定的分批量时继续收集数据
			if len(self.dataDocker) < self.dockerCap {
				continue
			}
			// 执行输出
			self.dataBatch++
			self.outputData()
		}
		// 将剩余收集到但未输出的数据输出
		self.dataBatch++
		self.outputData()
		close(dataStop)
	}()

	go func() {
		defer func() {
			recover()
		}()
		// 只有当收到退出通知并且通道内无数据时，才退出循环
		for file := range self.FileChan {
			atomic.AddUint64(&self.fileBatch, 1)
			self.wait.Add(1)
			go self.outputFile(file)
		}
		close(fileStop)
	}()

	<-dataStop
	<-fileStop

	// 等待所有输出完成
	self.wait.Wait()

	// 返回报告
	self.Report()
}

func (self *Collector) resetDataDocker() {
	for _, cell := range self.dataDocker {
		data.PutDataCell(cell)
	}
	self.dataDocker = self.dataDocker[:0]
}

// 获取文本数据总量
func (self *Collector) dataSum() uint64 {
	self.dataSumLock.RLock()
	defer self.dataSumLock.RUnlock()
	return self.sum[1]
}

// 更新文本数据总量
func (self *Collector) addDataSum(add uint64) {
	self.dataSumLock.Lock()
	defer self.dataSumLock.Unlock()
	self.sum[0] = self.sum[1]
	self.sum[1] += add
}

// 获取文件数据总量
func (self *Collector) fileSum() uint64 {
	self.fileSumLock.RLock()
	defer self.fileSumLock.RUnlock()
	return self.sum[3]
}

// 更新文件数据总量
func (self *Collector) addFileSum(add uint64) {
	self.fileSumLock.Lock()
	defer self.fileSumLock.Unlock()
	self.sum[2] = self.sum[3]
	self.sum[3] += add
}

// // 获取文本输出流量
// func (self *Collector) dataSize() uint64 {
// 	return self.size[0]
// }

// // 更新文本输出流量记录
// func (self *Collector) addDataSize(add uint64) {
// 	self.size[0] += add
// }

// // 获取文件输出流量
// func (self *Collector) fileSize() uint64 {
// 	return self.size[1]
// }

// // 更新文本输出流量记录
// func (self *Collector) addFileSize(add uint64) {
// 	self.size[1] += add
// }

// 返回报告
func (self *Collector) Report() {
	cache.ReportChan <- &cache.Report{
		SpiderName: self.Spider.GetName(),
		Keyin:      self.GetKeyin(),
		DataNum:    self.dataSum(),
		FileNum:    self.fileSum(),
		// DataSize:   self.dataSize(),
		// FileSize: self.fileSize(),
		Time: time.Since(cache.StartTime),
	}
}

func (self *Collector) Logger() logs.Logs {
	if self.logger == nil {
		return logs.Log
	}
	return self.logger
}

func (self *Collector) SetLogger(logger logs.Logs) {
	self.logger = logger
}
<commit_msg>update<commit_after>// 结果收集与输出
package collector

import (
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/admpub/spider/app/pipeline/collector/data"
	"github.com/admpub/spider/app/spider"
	"github.com/admpub/spider/config"
	"github.com/admpub/spider/logs"
	"github.com/admpub/spider/runtime/cache"
)

// 结果收集与输出
type Collector struct {
	*spider.Spider                    //绑定的采集规则
	DataChan       chan data.DataCell //文本数据收集通道
	FileChan       chan data.FileCell //文件收集通道
	dataDocker     []data.DataCell    //分批输出结果缓存
	outType        string             //输出方式
	fileOutType    string             //文件输出方式
	FileOutPath    string             //文件输出路径
	// size     [2]uint64 //数据总输出流量统计[文本，文件]，文本暂时未统计
	dataBatch   uint64 //当前文本输出批次
	fileBatch   uint64 //当前文件输出批次
	wait        sync.WaitGroup
	sum         [4]uint64 //收集的数据总数[上次输出后文本总数，本次输出后文本总数，上次输出后文件总数，本次输出后文件总数]，非并发安全
	dataSumLock sync.RWMutex
	fileSumLock sync.RWMutex
	logger      logs.Logs
	dockerCap   int
}

func NewCollector(sp *spider.Spider) *Collector {
	var self = &Collector{}
	self.Spider = sp
	if len(sp.OutType) > 0 {
		self.outType = sp.OutType
	} else {
		self.outType = cache.Task.OutType
	}
	if len(sp.FileOutType) > 0 {
		self.fileOutType = sp.FileOutType
	} else {
		self.fileOutType = cache.Task.FileOutType
	}
	if len(self.fileOutType) == 0 {
		self.fileOutType = `local`
	}
	if len(sp.FileOutPath) > 0 {
		self.FileOutPath = sp.FileOutPath
	} else {
		self.FileOutPath = cache.Task.FileOutPath
	}
	if len(self.FileOutPath) == 0 {
		self.FileOutPath = config.FILE_DIR
	}
	if sp.DockerCap > 0 {
		self.dockerCap = sp.DockerCap
	} else {
		if cache.Task.DockerCap < 1 {
			cache.Task.DockerCap = 1
		}
		self.dockerCap = cache.Task.DockerCap
	}
	self.DataChan = make(chan data.DataCell, self.dockerCap)
	self.FileChan = make(chan data.FileCell, self.dockerCap)
	self.dataDocker = make([]data.DataCell, 0, self.dockerCap)
	self.sum = [4]uint64{}
	// self.size = [2]uint64{}
	self.dataBatch = 0
	self.fileBatch = 0
	return self
}

func (self *Collector) CollectData(dataCell data.DataCell) error {
	var err error
	defer func() {
		if recover() != nil {
			err = fmt.Errorf("输出协程已终止")
		}
	}()
	self.DataChan <- dataCell
	return err
}

func (self *Collector) CollectFile(fileCell data.FileCell) error {
	var err error
	defer func() {
		if recover() != nil {
			err = fmt.Errorf("输出协程已终止")
		}
	}()
	self.FileChan <- fileCell
	return err
}

// 停止
func (self *Collector) Stop() {
	go func() {
		defer func() {
			recover()
		}()
		close(self.DataChan)
	}()
	go func() {
		defer func() {
			recover()
		}()
		close(self.FileChan)
	}()
}

// 启动数据收集/输出管道
func (self *Collector) Start() {
	// 启动输出协程
	go self.Run()
}

// 启动数据收集/输出管道
func (self *Collector) Run() {
	dataStop := make(chan bool)
	fileStop := make(chan bool)

	go func() {
		defer func() {
			recover()
		}()
		for data := range self.DataChan {
			// 缓存分批数据
			self.dataDocker = append(self.dataDocker, data)

			// 未达到设定的分批量时继续收集数据
			if len(self.dataDocker) < self.dockerCap {
				continue
			}
			// 执行输出
			self.dataBatch++
			self.outputData()
		}
		// 将剩余收集到但未输出的数据输出
		self.dataBatch++
		self.outputData()
		close(dataStop)
	}()

	go func() {
		defer func() {
			recover()
		}()
		// 只有当收到退出通知并且通道内无数据时，才退出循环
		for file := range self.FileChan {
			atomic.AddUint64(&self.fileBatch, 1)
			self.wait.Add(1)
			go self.outputFile(file)
		}
		close(fileStop)
	}()

	<-dataStop
	<-fileStop

	// 等待所有输出完成
	self.wait.Wait()

	// 返回报告
	self.Report()
}

func (self *Collector) resetDataDocker() {
	for _, cell := range self.dataDocker {
		data.PutDataCell(cell)
	}
	self.dataDocker = self.dataDocker[:0]
}

// 获取文本数据总量
func (self *Collector) dataSum() uint64 {
	self.dataSumLock.RLock()
	defer self.dataSumLock.RUnlock()
	return self.sum[1]
}

// 更新文本数据总量
func (self *Collector) addDataSum(add uint64) {
	self.dataSumLock.Lock()
	defer self.dataSumLock.Unlock()
	self.sum[0] = self.sum[1]
	self.sum[1] += add
}

// 获取文件数据总量
func (self *Collector) fileSum() uint64 {
	self.fileSumLock.RLock()
	defer self.fileSumLock.RUnlock()
	return self.sum[3]
}

// 更新文件数据总量
func (self *Collector) addFileSum(add uint64) {
	self.fileSumLock.Lock()
	defer self.fileSumLock.Unlock()
	self.sum[2] = self.sum[3]
	self.sum[3] += add
}

// // 获取文本输出流量
// func (self *Collector) dataSize() uint64 {
// 	return self.size[0]
// }

// // 更新文本输出流量记录
// func (self *Collector) addDataSize(add uint64) {
// 	self.size[0] += add
// }

// // 获取文件输出流量
// func (self *Collector) fileSize() uint64 {
// 	return self.size[1]
// }

// // 更新文本输出流量记录
// func (self *Collector) addFileSize(add uint64) {
// 	self.size[1] += add
// }

// 返回报告
func (self *Collector) Report() {
	cache.ReportChan <- &cache.Report{
		SpiderName: self.Spider.GetName(),
		Keyin:      self.GetKeyin(),
		DataNum:    self.dataSum(),
		FileNum:    self.fileSum(),
		// DataSize:   self.dataSize(),
		// FileSize: self.fileSize(),
		Time: time.Since(cache.StartTime),
	}
}

func (self *Collector) Logger() logs.Logs {
	if self.logger == nil {
		return logs.Log
	}
	return self.logger
}

func (self *Collector) SetLogger(logger logs.Logs) {
	self.logger = logger
}
<|endoftext|>