<commit_before>package rep

import (
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/cloudfoundry-incubator/rep"
	"github.com/onsi/say"
)

func RepState(out io.Writer) (err error) {
	client := rep.NewClient(&http.Client{
		Timeout: 5 * time.Second,
	}, "http://localhost:1800")

	t := time.Now()
	state, err := client.State()
	dt := time.Since(t)

	if err != nil {
		say.Println(0, "Cell State [%s] - Error:%s", dt, say.Red(err.Error()))
		return err
	}

	name := say.Green("Cell State")
	if state.Evacuating {
		name = say.Red("Cell State - EVAC -")
	}

	rootFSes := []string{}
	for key := range state.RootFSProviders {
		if key != "preloaded" {
			rootFSes = append(rootFSes, say.Yellow(key))
		}
	}

	for key := range state.RootFSProviders["preloaded"].(rep.FixedSetRootFSProvider).FixedSet {
		rootFSes = append(rootFSes, say.Green("preloaded:%s", key))
	}

	say.Println(0, "%s [%s] - Zone:%s | %s Tasks, %s LRPs | C:%d/%d M:%d/%d D:%d/%d | %s",
		name,
		dt,
		say.Cyan(state.Zone),
		say.Cyan("%d", len(state.Tasks)),
		say.Cyan("%d", len(state.LRPs)),
		state.AvailableResources.Containers,
		state.TotalResources.Containers,
		state.AvailableResources.MemoryMB,
		state.TotalResources.MemoryMB,
		state.AvailableResources.DiskMB,
		state.TotalResources.DiskMB,
		strings.Join(rootFSes, ", "),
	)

	return nil
}
<commit_msg>Update veritas to pass second http client when creating rep client<commit_after>package rep

import (
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/cloudfoundry-incubator/rep"
	"github.com/onsi/say"
)

func RepState(out io.Writer) (err error) {
	httpClient := &http.Client{
		Timeout: 5 * time.Second,
	}
	client := rep.NewClient(httpClient, httpClient, "http://localhost:1800")

	t := time.Now()
	state, err := client.State()
	dt := time.Since(t)

	if err != nil {
		say.Println(0, "Cell State [%s] - Error:%s", dt, say.Red(err.Error()))
		return err
	}

	name := say.Green("Cell State")
	if state.Evacuating {
		name = say.Red("Cell State - EVAC -")
	}

	rootFSes := []string{}
	for key := range state.RootFSProviders {
		if key != "preloaded" {
			rootFSes = append(rootFSes, say.Yellow(key))
		}
	}

	for key := range state.RootFSProviders["preloaded"].(rep.FixedSetRootFSProvider).FixedSet {
		rootFSes = append(rootFSes, say.Green("preloaded:%s", key))
	}

	say.Println(0, "%s [%s] - Zone:%s | %s Tasks, %s LRPs | C:%d/%d M:%d/%d D:%d/%d | %s",
		name,
		dt,
		say.Cyan(state.Zone),
		say.Cyan("%d", len(state.Tasks)),
		say.Cyan("%d", len(state.LRPs)),
		state.AvailableResources.Containers,
		state.TotalResources.Containers,
		state.AvailableResources.MemoryMB,
		state.TotalResources.MemoryMB,
		state.AvailableResources.DiskMB,
		state.TotalResources.DiskMB,
		strings.Join(rootFSes, ", "),
	)

	return nil
}
<|endoftext|>


<commit_before>package logpeck

import (
	"github.com/Sirupsen/logrus"
	log "github.com/Sirupsen/logrus"
	"strconv"
	"time"
)

type AggregatorConfig struct {
	Tags         []string `json:"Tags"`
	Aggregations []string `json:"Aggregations"`
	Target       string   `json:"Target"`
	PreFields    string   `json:"PreFields"`
	Timestamp    string   `json:"Timestamp"`
}

type Aggregator struct {
	Interval          int64
	FieldsKey         string
	AggregatorConfigs map[string]AggregatorConfig
	buckets           map[string]map[string][]int64
	postTime          int64
}

func NewAggregator(interval int64, fieldsKey string, aggregators *map[string]AggregatorConfig) *Aggregator {
	aggregator := &Aggregator{
		Interval:          interval,
		FieldsKey:         fieldsKey,
		AggregatorConfigs: *aggregators,
		buckets:           make(map[string]map[string][]int64),
		postTime:          0,
	}
	return aggregator
}

func getSampleTime(ts int64, interval int64) int64 {
	return ts / interval
}

func (p *Aggregator) IsDeadline(timestamp int64) bool {
	interval := p.Interval
	nowTime := getSampleTime(timestamp, interval)
	if p.postTime != nowTime {
		return true
	}
	return false
}

func (p *Aggregator) Record(fields map[string]interface{}) int64 {
	bucketName := fields[p.FieldsKey].(string)
	bucketTag := ""
	aggregatorConfig := p.AggregatorConfigs[bucketName]
	tags := aggregatorConfig.Tags
	aggregations := aggregatorConfig.Aggregations
	target := aggregatorConfig.Target
	timestamp := aggregatorConfig.Timestamp
	preFields := aggregatorConfig.PreFields
	if target == "" {
		return time.Now().Unix()
	}
	for i := 0; i < len(tags); i++ {
		bucketTag += "," + tags[i] + "=" + fields[tags[i]].(string)
	}
	if preFields == "" {
		bucketTag += " "
	} else {
		bucketTag += " " + fields[preFields].(string) + "_"
	}
	int_bool := false
	for i := 0; i < len(aggregations); i++ {
		if aggregations[i] != "cnt" {
			int_bool = true
		}
	}

	aggValue := fields[target].(string)

	//get time
	now, err := strconv.ParseInt(fields[timestamp].(string), 10, 64)
	if err != nil {
		logrus.Infof("[Record] timestamp:%v can't use strconv.ParseInt", fields[timestamp].(string))
		now = time.Now().Unix()
	}

	if _, ok := p.buckets[bucketName]; !ok {
		p.buckets[bucketName] = make(map[string][]int64)
	}
	if int_bool == false {
		p.buckets[bucketName][bucketTag] = append(p.buckets[bucketName][bucketTag], 1)
	} else {
		aggValue, err := strconv.ParseInt(aggValue, 10, 64)
		if err != nil {
			logrus.Infof("[Record] target:%v can't use strconv.ParseInt", aggValue)
			return now
		}
		p.buckets[bucketName][bucketTag] = append(p.buckets[bucketName][bucketTag], aggValue)
	}

	return now
}

func quickSort(values []int64, left, right int64) {
	temp := values[left]
	p := left
	i, j := left, right
	for i <= j {
		for j >= p && values[j] >= temp {
			j--
		}
		if j >= p {
			values[p] = values[j]
			p = j
		}
		for i <= p && values[i] <= temp {
			i++
		}
		if i <= p {
			values[p] = values[i]
			p = i
		}
	}
	values[p] = temp

	if p-left > 1 {
		quickSort(values, left, p-1)
	}
	if right-p > 1 {
		quickSort(values, p+1, right)
	}
}

func getAggregation(targetValue []int64, aggregations []string) map[string]int64 {
	aggregationResults := map[string]int64{}
	cnt := int64(len(targetValue))
	avg := int64(0)
	sum := int64(0)
	min := int64(0)
	max := int64(0)
	if cnt > 0 {
		min = targetValue[0]
		max = targetValue[0]
	}
	quickSort(targetValue, int64(0), int64(len(targetValue)-1))
	for _, value := range targetValue {
		sum += value
		if value > max {
			max = value
		}
		if value < min {
			min = value
		}
	}
	avg = sum / cnt
	for i := 0; i < len(aggregations); i++ {
		switch aggregations[i] {
		case "cnt":
			aggregationResults["cnt"] = int64(len(targetValue))
		case "sum":
			aggregationResults["sum"] = sum
		case "avg":
			aggregationResults["avg"] = avg
		case "min":
			aggregationResults["min"] = min
		case "max":
			aggregationResults["max"] = max
		default:
			if aggregations[i][0] == 'p' {
				proportion, err := strconv.ParseInt(aggregations[i][1:], 10, 64)
				if err != nil {
					panic(aggregations[i])
				}
				percentile := targetValue[cnt*proportion/100-1]
				aggregationResults[aggregations[i]] = percentile
			}
		}
	}
	return aggregationResults
}

func (p *Aggregator) Dump(timestamp int64) map[string]interface{} {
	fields := map[string]interface{}{}
	//now := strconv.FormatInt(timestamp, 10)
	for bucketName, bucketTag_value := range p.buckets {
		for bucketTag, targetValue := range bucketTag_value {
			aggregations := p.AggregatorConfigs[bucketName].Aggregations
			fields[bucketName+bucketTag] = getAggregation(targetValue, aggregations)
		}
	}
	fields["timestamp"] = timestamp
	p.postTime = getSampleTime(timestamp, p.Interval)
	p.buckets = map[string]map[string][]int64{}
	log.Infof("[Dump] fields is : %v", fields)
	return fields
}
<commit_msg>Modify the error of type conversion<commit_after>package logpeck

import (
	log "github.com/Sirupsen/logrus"
	"strconv"
	"time"
)

type AggregatorConfig struct {
	Tags         []string `json:"Tags"`
	Aggregations []string `json:"Aggregations"`
	Target       string   `json:"Target"`
	PreFields    string   `json:"PreFields"`
	Timestamp    string   `json:"Timestamp"`
}

type Aggregator struct {
	Interval          int64
	FieldsKey         string
	AggregatorConfigs map[string]AggregatorConfig
	buckets           map[string]map[string][]int64
	postTime          int64
}

func NewAggregator(interval int64, fieldsKey string, aggregators *map[string]AggregatorConfig) *Aggregator {
	aggregator := &Aggregator{
		Interval:          interval,
		FieldsKey:         fieldsKey,
		AggregatorConfigs: *aggregators,
		buckets:           make(map[string]map[string][]int64),
		postTime:          0,
	}
	return aggregator
}

func getSampleTime(ts int64, interval int64) int64 {
	return ts / interval
}

func (p *Aggregator) IsDeadline(timestamp int64) bool {
	interval := p.Interval
	nowTime := getSampleTime(timestamp, interval)
	if p.postTime != nowTime {
		return true
	}
	return false
}

func (p *Aggregator) Record(fields map[string]interface{}) int64 {
	bucketName := fields[p.FieldsKey].(string)
	bucketTag := ""
	aggregatorConfig := p.AggregatorConfigs[bucketName]
	tags := aggregatorConfig.Tags
	target := aggregatorConfig.Target
	timestamp := aggregatorConfig.Timestamp
	preFields := aggregatorConfig.PreFields
	if target == "" {
		return time.Now().Unix()
	}
	for i := 0; i < len(tags); i++ {
		bucketTag += "," + tags[i] + "=" + fields[tags[i]].(string)
	}
	if preFields == "" {
		bucketTag += " "
	} else {
		bucketTag += " " + fields[preFields].(string) + "_"
	}
	aggValue := fields[target].(string)

	//get time
	now, err := strconv.ParseInt(fields[timestamp].(string), 10, 64)
	if err != nil {
		log.Debug("[Record] timestamp:%v can't use strconv.ParseInt", fields[timestamp].(string))
		now = time.Now().Unix()
	}

	if _, ok := p.buckets[bucketName]; !ok {
		p.buckets[bucketName] = make(map[string][]int64)
	}
	aggValueInt, err := strconv.ParseInt(aggValue, 10, 64)
	if err != nil {
		log.Infof("[Record] target:%v can't use strconv.ParseInt", aggValue)
		p.buckets[bucketName][bucketTag] = append(p.buckets[bucketName][bucketTag], 1)
	} else {
		p.buckets[bucketName][bucketTag] = append(p.buckets[bucketName][bucketTag], aggValueInt)
	}
	return now
}

func quickSort(values []int64, left, right int64) {
	temp := values[left]
	p := left
	i, j := left, right
	for i <= j {
		for j >= p && values[j] >= temp {
			j--
		}
		if j >= p {
			values[p] = values[j]
			p = j
		}
		for i <= p && values[i] <= temp {
			i++
		}
		if i <= p {
			values[p] = values[i]
			p = i
		}
	}
	values[p] = temp

	if p-left > 1 {
		quickSort(values, left, p-1)
	}
	if right-p > 1 {
		quickSort(values, p+1, right)
	}
}

func getAggregation(targetValue []int64, aggregations []string) map[string]int64 {
	log.Infof("[getAggregation] targetValue is : %v", targetValue)
	aggregationResults := map[string]int64{}
	cnt := int64(len(targetValue))
	avg := int64(0)
	sum := int64(0)
	min := int64(0)
	max := int64(0)
	if cnt > 0 {
		min = targetValue[0]
		max = targetValue[0]
	}
	quickSort(targetValue, int64(0), int64(len(targetValue)-1))
	for _, value := range targetValue {
		sum += value
		if value > max {
			max = value
		}
		if value < min {
			min = value
		}
	}
	avg = sum / cnt
	for i := 0; i < len(aggregations); i++ {
		switch aggregations[i] {
		case "cnt":
			aggregationResults["cnt"] = int64(len(targetValue))
		case "sum":
			aggregationResults["sum"] = sum
		case "avg":
			aggregationResults["avg"] = avg
		case "min":
			aggregationResults["min"] = min
		case "max":
			aggregationResults["max"] = max
		default:
			if aggregations[i][0] == 'p' {
				proportion, err := strconv.ParseInt(aggregations[i][1:], 10, 64)
				if err != nil {
					panic(aggregations[i])
				}
				index := cnt*proportion/100 - 1
				if cnt*proportion/100-1 < 0 {
					index = 0
				}
				percentile := targetValue[index]
				aggregationResults[aggregations[i]] = percentile
			}
		}
	}
	return aggregationResults
}

func (p *Aggregator) Dump(timestamp int64) map[string]interface{} {
	fields := map[string]interface{}{}
	log.Infof("[Dump] bucket is : %v", p.buckets)
	//now := strconv.FormatInt(timestamp, 10)
	for bucketName, bucketTag_value := range p.buckets {
		for bucketTag, targetValue := range bucketTag_value {
			aggregations := p.AggregatorConfigs[bucketName].Aggregations
			fields[bucketName+bucketTag] = getAggregation(targetValue, aggregations)
		}
	}
	fields["timestamp"] = timestamp
	p.postTime = getSampleTime(timestamp, p.Interval)
	p.buckets = map[string]map[string][]int64{}
	log.Infof("[Dump] fields is : %v", fields)
	return fields
}
<|endoftext|>


<commit_before>// example to demonstrate reading query parameters (with multiple values in the below format) using net/http library.
// example input: https://localhost:8080/customer?country=usa&country=india
// required output: a slice of string with values usa and india like this []string{"usa","india"}
//

package main

import (
	"fmt"
	"net/http"
	"net/url"
)

func queryParamHandler(w http.ResponseWriter, r *http.Request) {
	vars, err := url.ParseQuery(r.URL.RawQuery)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
	}
	var country []string
	country = vars["country"]

	fmt.Fprintln(w, country)
}

func main() {

	http.HandleFunc("/customer", queryParamHandler)

	err := http.ListenAndServe(":8080", nil)
	if err != nil {
		panic(err)
	}
}
<commit_msg>Update read-query-param-multiple2.go<commit_after>// example to demonstrate reading query parameters (with multiple values in the below format) using net/http library.
// example input: http://localhost:8080/customer?country=usa&country=india
// required output: a slice of string with values usa and india like this []string{"usa","india"}
//

package main

import (
	"fmt"
	"net/http"
	"net/url"
)

func queryParamHandler(w http.ResponseWriter, r *http.Request) {
	vars, err := url.ParseQuery(r.URL.RawQuery)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
	}
	var country []string
	country = vars["country"]

	fmt.Fprintln(w, country)
}

func main() {

	http.HandleFunc("/customer", queryParamHandler)

	err := http.ListenAndServe(":8080", nil)
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 Jigsaw Operations LLC
//
// 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
//
//     https://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 service

import (
	"bytes"
	"container/list"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"sync"
	"time"

	onet "github.com/Jigsaw-Code/outline-ss-server/net"
	"github.com/Jigsaw-Code/outline-ss-server/service/metrics"
	ss "github.com/Jigsaw-Code/outline-ss-server/shadowsocks"
	logging "github.com/op/go-logging"
	"github.com/shadowsocks/go-shadowsocks2/socks"
)

func remoteIP(conn net.Conn) net.IP {
	addr := conn.RemoteAddr()
	if addr == nil {
		return nil
	}
	if tcpaddr, ok := addr.(*net.TCPAddr); ok {
		return tcpaddr.IP
	}
	ipstr, _, err := net.SplitHostPort(addr.String())
	if err == nil {
		return net.ParseIP(ipstr)
	}
	return nil
}

// Wrapper for logger.Debugf during TCP access key searches.
func debugTCP(cipherID, template string, val interface{}) {
	// This is an optimization to reduce unnecessary allocations due to an interaction
	// between Go's inlining/escape analysis and varargs functions like logger.Debugf.
	if logger.IsEnabledFor(logging.DEBUG) {
		logger.Debugf("TCP(%s): "+template, cipherID, val)
	}
}

func findAccessKey(clientReader io.Reader, clientIP net.IP, cipherList CipherList) (*CipherEntry, io.Reader, []byte, time.Duration, error) {
	// We snapshot the list because it may be modified while we use it.
	tcpTrialSize, ciphers := cipherList.SnapshotForClientIP(clientIP)
	firstBytes := make([]byte, tcpTrialSize)
	if n, err := io.ReadFull(clientReader, firstBytes); err != nil {
		return nil, clientReader, nil, 0, fmt.Errorf("Reading header failed after %d bytes: %v", n, err)
	}

	findStartTime := time.Now()
	entry, elt := findEntry(firstBytes, ciphers)
	timeToCipher := time.Now().Sub(findStartTime)
	if entry == nil {
		// TODO: Ban and log client IPs with too many failures too quick to protect against DoS.
		return nil, clientReader, nil, timeToCipher, fmt.Errorf("Could not find valid TCP cipher")
	}

	// Move the active cipher to the front, so that the search is quicker next time.
	cipherList.MarkUsedByClientIP(elt, clientIP)
	salt := firstBytes[:entry.Cipher.SaltSize()]
	return entry, io.MultiReader(bytes.NewReader(firstBytes), clientReader), salt, timeToCipher, nil
}

// Implements a trial decryption search.  This assumes that all ciphers are AEAD.
func findEntry(firstBytes []byte, ciphers []*list.Element) (*CipherEntry, *list.Element) {
	// Constant of zeroes to use as the start chunk count.
	zeroCountBuf := [maxNonceSize]byte{}
	// To hold the decrypted chunk length.
	chunkLenBuf := [2]byte{}
	for ci, elt := range ciphers {
		entry := elt.Value.(*CipherEntry)
		id, cipher := entry.ID, entry.Cipher
		saltsize := cipher.SaltSize()
		salt := firstBytes[:saltsize]
		aead, err := cipher.Decrypter(salt)
		if err != nil {
			debugTCP(id, "Failed to create decrypter: %v", err)
			continue
		}
		cipherTextLength := 2 + aead.Overhead()
		cipherText := firstBytes[saltsize : saltsize+cipherTextLength]
		_, err = aead.Open(chunkLenBuf[:0], zeroCountBuf[:aead.NonceSize()], cipherText, nil)
		if err != nil {
			debugTCP(id, "Failed to decrypt length: %v", err)
			continue
		}
		debugTCP(id, "Found cipher at index %d", ci)
		// Move the active cipher to the front, so that the search is quicker next time.
		return entry, elt
	}
	return nil, nil
}

type tcpService struct {
	mu          sync.RWMutex // Protects .listeners and .stopped
	listener    *net.TCPListener
	stopped     bool
	ciphers     CipherList
	m           metrics.ShadowsocksMetrics
	running     sync.WaitGroup
	readTimeout time.Duration
	// `replayCache` is a pointer to SSServer.replayCache, to share the cache among all ports.
	replayCache       *ReplayCache
	targetIPValidator onet.TargetIPValidator
}

// NewTCPService creates a TCPService
// `replayCache` is a pointer to SSServer.replayCache, to share the cache among all ports.
func NewTCPService(ciphers CipherList, replayCache *ReplayCache, m metrics.ShadowsocksMetrics, timeout time.Duration) TCPService {
	return &tcpService{
		ciphers:           ciphers,
		m:                 m,
		readTimeout:       timeout,
		replayCache:       replayCache,
		targetIPValidator: onet.RequirePublicIP,
	}
}

// TCPService is a Shadowsocks TCP service that can be started and stopped.
type TCPService interface {
	// SetTargetIPValidator sets the function to be used to validate the target IP addresses.
	SetTargetIPValidator(targetIPValidator onet.TargetIPValidator)
	// Serve adopts the listener, which will be closed before Serve returns.  Serve returns an error unless Stop() was called.
	Serve(listener *net.TCPListener) error
	// Stop closes the listener but does not interfere with existing connections.
	Stop() error
	// GracefulStop calls Stop(), and then blocks until all resources have been cleaned up.
	GracefulStop() error
}

func (s *tcpService) SetTargetIPValidator(targetIPValidator onet.TargetIPValidator) {
	s.targetIPValidator = targetIPValidator
}

// proxyConnection will route the clientConn according to the address read from the connection.
func proxyConnection(clientSSConn onet.DuplexConn, tgtAddr socks.Addr, proxyMetrics *metrics.ProxyMetrics, targetIPValidator onet.TargetIPValidator) *onet.ConnectionError {
	tgtTCPAddr, err := net.ResolveTCPAddr("tcp", tgtAddr.String())
	if err != nil {
		return onet.NewConnectionError("ERR_RESOLVE_ADDRESS", fmt.Sprintf("Failed to resolve target address %v", tgtAddr.String()), err)
	}
	if err := targetIPValidator(tgtTCPAddr.IP); err != nil {
		return err
	}

	tgtTCPConn, err := net.DialTCP("tcp", nil, tgtTCPAddr)
	if err != nil {
		return onet.NewConnectionError("ERR_CONNECT", "Failed to connect to target", err)
	}
	defer tgtTCPConn.Close()
	tgtTCPConn.SetKeepAlive(true)
	tgtConn := metrics.MeasureConn(tgtTCPConn, &proxyMetrics.ProxyTarget, &proxyMetrics.TargetProxy)

	logger.Debugf("proxy %s <-> %s", clientSSConn.RemoteAddr().String(), tgtConn.RemoteAddr().String())
	fromClientErrCh := make(chan error)
	go func() {
		_, fromClientErr := clientSSConn.(io.WriterTo).WriteTo(tgtConn)
		// Send FIN to target.
		tgtConn.CloseWrite()
		if fromClientErr != nil {
			// Drain to prevent a close on cipher error.
			// TODO: Need to drain the underlying connections instead.
			clientSSConn.(io.WriterTo).WriteTo(ioutil.Discard)
		}
		clientSSConn.CloseRead()
		fromClientErrCh <- fromClientErr
	}()
	_, fromTargetErr := clientSSConn.(io.ReaderFrom).ReadFrom(tgtConn)
	// Send FIN to client.
	clientSSConn.CloseWrite()
	tgtConn.CloseRead()

	fromClientErr := <-fromClientErrCh
	if fromClientErr != nil {
		return onet.NewConnectionError("ERR_RELAY_CLIENT", "Failed to relay traffic from client", fromClientErr)
	}
	if fromTargetErr != nil {
		return onet.NewConnectionError("ERR_RELAY_TARGET", "Failed to relay traffic from target", fromTargetErr)
	}
	return nil
}

func (s *tcpService) Serve(listener *net.TCPListener) error {
	s.mu.Lock()
	if s.listener != nil {
		s.mu.Unlock()
		listener.Close()
		return errors.New("Serve can only be called once")
	}
	if s.stopped {
		s.mu.Unlock()
		return listener.Close()
	}
	s.listener = listener
	s.running.Add(1)
	s.mu.Unlock()

	defer s.running.Done()
	for {
		clientConn, err := listener.AcceptTCP()
		if err != nil {
			s.mu.RLock()
			stopped := s.stopped
			s.mu.RUnlock()
			if stopped {
				return nil
			}
			logger.Errorf("Accept failed: %v", err)
			continue
		}

		s.running.Add(1)
		go func() {
			defer s.running.Done()
			defer func() {
				if r := recover(); r != nil {
					logger.Errorf("Panic in TCP handler: %v", r)
				}
			}()
			s.handleConnection(listener.Addr().(*net.TCPAddr).Port, clientConn)
		}()
	}
}

func (s *tcpService) handleConnection(listenerPort int, clientTCPConn *net.TCPConn) {
	clientLocation, err := s.m.GetLocation(clientTCPConn.RemoteAddr())
	if err != nil {
		logger.Warningf("Failed location lookup: %v", err)
	}
	logger.Debugf("Got location \"%v\" for IP %v", clientLocation, clientTCPConn.RemoteAddr().String())
	s.m.AddOpenTCPConnection(clientLocation)

	connStart := time.Now()
	clientTCPConn.SetKeepAlive(true)
	// Set a deadline to receive the address to the target.
	clientTCPConn.SetReadDeadline(connStart.Add(s.readTimeout))
	var proxyMetrics metrics.ProxyMetrics
	clientConn := metrics.MeasureConn(clientTCPConn, &proxyMetrics.ProxyClient, &proxyMetrics.ClientProxy)
	cipherEntry, clientReader, clientSalt, timeToCipher, keyErr := findAccessKey(clientConn, remoteIP(clientConn), s.ciphers)

	connError := func() *onet.ConnectionError {
		if keyErr != nil {
			logger.Debugf("Failed to find a valid cipher after reading %v bytes: %v", proxyMetrics.ClientProxy, keyErr)
			const status = "ERR_CIPHER"
			s.absorbProbe(listenerPort, clientConn, clientLocation, status, &proxyMetrics)
			return onet.NewConnectionError(status, "Failed to find a valid cipher", keyErr)
		}

		isServerSalt := cipherEntry.SaltGenerator.IsServerSalt(clientSalt)
		// Only check the cache if findAccessKey succeeded and the salt is unrecognized.
		if isServerSalt || !s.replayCache.Add(cipherEntry.ID, clientSalt) {
			var status string
			if isServerSalt {
				status = "ERR_REPLAY_SERVER"
			} else {
				status = "ERR_REPLAY_CLIENT"
			}
			s.absorbProbe(listenerPort, clientConn, clientLocation, status, &proxyMetrics)
			logger.Debugf(status+": %v in %s sent %d bytes", clientConn.RemoteAddr(), clientLocation, proxyMetrics.ClientProxy)
			return onet.NewConnectionError(status, "Replay detected", nil)
		}

		ssr := ss.NewShadowsocksReader(clientReader, cipherEntry.Cipher)
		tgtAddr, err := socks.ReadAddr(ssr)
		if err != nil {
			// Drain to prevent a close on cipher error.
			clientConn.(io.WriterTo).WriteTo(ioutil.Discard)
			return onet.NewConnectionError("ERR_READ_ADDRESS", "Failed to get target address", err)
		}

		// Clear the deadline for the target address
		clientConn.SetReadDeadline(time.Time{})
		ssw := ss.NewShadowsocksWriter(clientConn, cipherEntry.Cipher)
		ssw.SetSaltGenerator(cipherEntry.SaltGenerator)
		clientSSConn := onet.WrapConn(clientConn, ssr, ssw)
		return proxyConnection(clientSSConn, tgtAddr, &proxyMetrics, s.targetIPValidator)
	}()

	connDuration := time.Now().Sub(connStart)
	status := "OK"
	if connError != nil {
		logger.Debugf("TCP Error: %v: %v", connError.Message, connError.Cause)
		status = connError.Status
	}
	var id string
	if cipherEntry != nil {
		id = cipherEntry.ID
	}
	s.m.AddClosedTCPConnection(clientLocation, id, status, proxyMetrics, timeToCipher, connDuration)
	clientConn.Close() // Closing after the metrics are added aids integration testing.
	logger.Debugf("Done with status %v, duration %v", status, connDuration)
}

// Keep the connection open until we hit the authentication deadline to protect against probing attacks
// `proxyMetrics` is a pointer because its value is being mutated by `clientConn`.
func (s *tcpService) absorbProbe(listenerPort int, clientConn io.ReadCloser, clientLocation, status string, proxyMetrics *metrics.ProxyMetrics) {
	_, drainErr := io.Copy(ioutil.Discard, clientConn) // drain socket
	drainResult := drainErrToString(drainErr)
	logger.Debugf("Drain error: %v, drain result: %v", drainErr, drainResult)
	s.m.AddTCPProbe(clientLocation, status, drainResult, listenerPort, *proxyMetrics)
}

func drainErrToString(drainErr error) string {
	netErr, ok := drainErr.(net.Error)
	switch {
	case drainErr == nil:
		return "eof"
	case ok && netErr.Timeout():
		return "timeout"
	default:
		return "other"
	}
}

func (s *tcpService) Stop() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.stopped = true
	if s.listener == nil {
		return nil
	}
	return s.listener.Close()
}

func (s *tcpService) GracefulStop() error {
	err := s.Stop()
	s.running.Wait()
	return err
}
<commit_msg>Move tgtConn.CloseWrite<commit_after>// Copyright 2018 Jigsaw Operations LLC
//
// 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
//
//     https://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 service

import (
	"bytes"
	"container/list"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"sync"
	"time"

	onet "github.com/Jigsaw-Code/outline-ss-server/net"
	"github.com/Jigsaw-Code/outline-ss-server/service/metrics"
	ss "github.com/Jigsaw-Code/outline-ss-server/shadowsocks"
	logging "github.com/op/go-logging"
	"github.com/shadowsocks/go-shadowsocks2/socks"
)

func remoteIP(conn net.Conn) net.IP {
	addr := conn.RemoteAddr()
	if addr == nil {
		return nil
	}
	if tcpaddr, ok := addr.(*net.TCPAddr); ok {
		return tcpaddr.IP
	}
	ipstr, _, err := net.SplitHostPort(addr.String())
	if err == nil {
		return net.ParseIP(ipstr)
	}
	return nil
}

// Wrapper for logger.Debugf during TCP access key searches.
func debugTCP(cipherID, template string, val interface{}) {
	// This is an optimization to reduce unnecessary allocations due to an interaction
	// between Go's inlining/escape analysis and varargs functions like logger.Debugf.
	if logger.IsEnabledFor(logging.DEBUG) {
		logger.Debugf("TCP(%s): "+template, cipherID, val)
	}
}

func findAccessKey(clientReader io.Reader, clientIP net.IP, cipherList CipherList) (*CipherEntry, io.Reader, []byte, time.Duration, error) {
	// We snapshot the list because it may be modified while we use it.
	tcpTrialSize, ciphers := cipherList.SnapshotForClientIP(clientIP)
	firstBytes := make([]byte, tcpTrialSize)
	if n, err := io.ReadFull(clientReader, firstBytes); err != nil {
		return nil, clientReader, nil, 0, fmt.Errorf("Reading header failed after %d bytes: %v", n, err)
	}

	findStartTime := time.Now()
	entry, elt := findEntry(firstBytes, ciphers)
	timeToCipher := time.Now().Sub(findStartTime)
	if entry == nil {
		// TODO: Ban and log client IPs with too many failures too quick to protect against DoS.
		return nil, clientReader, nil, timeToCipher, fmt.Errorf("Could not find valid TCP cipher")
	}

	// Move the active cipher to the front, so that the search is quicker next time.
	cipherList.MarkUsedByClientIP(elt, clientIP)
	salt := firstBytes[:entry.Cipher.SaltSize()]
	return entry, io.MultiReader(bytes.NewReader(firstBytes), clientReader), salt, timeToCipher, nil
}

// Implements a trial decryption search.  This assumes that all ciphers are AEAD.
func findEntry(firstBytes []byte, ciphers []*list.Element) (*CipherEntry, *list.Element) {
	// Constant of zeroes to use as the start chunk count.
	zeroCountBuf := [maxNonceSize]byte{}
	// To hold the decrypted chunk length.
	chunkLenBuf := [2]byte{}
	for ci, elt := range ciphers {
		entry := elt.Value.(*CipherEntry)
		id, cipher := entry.ID, entry.Cipher
		saltsize := cipher.SaltSize()
		salt := firstBytes[:saltsize]
		aead, err := cipher.Decrypter(salt)
		if err != nil {
			debugTCP(id, "Failed to create decrypter: %v", err)
			continue
		}
		cipherTextLength := 2 + aead.Overhead()
		cipherText := firstBytes[saltsize : saltsize+cipherTextLength]
		_, err = aead.Open(chunkLenBuf[:0], zeroCountBuf[:aead.NonceSize()], cipherText, nil)
		if err != nil {
			debugTCP(id, "Failed to decrypt length: %v", err)
			continue
		}
		debugTCP(id, "Found cipher at index %d", ci)
		// Move the active cipher to the front, so that the search is quicker next time.
		return entry, elt
	}
	return nil, nil
}

type tcpService struct {
	mu          sync.RWMutex // Protects .listeners and .stopped
	listener    *net.TCPListener
	stopped     bool
	ciphers     CipherList
	m           metrics.ShadowsocksMetrics
	running     sync.WaitGroup
	readTimeout time.Duration
	// `replayCache` is a pointer to SSServer.replayCache, to share the cache among all ports.
	replayCache       *ReplayCache
	targetIPValidator onet.TargetIPValidator
}

// NewTCPService creates a TCPService
// `replayCache` is a pointer to SSServer.replayCache, to share the cache among all ports.
func NewTCPService(ciphers CipherList, replayCache *ReplayCache, m metrics.ShadowsocksMetrics, timeout time.Duration) TCPService {
	return &tcpService{
		ciphers:           ciphers,
		m:                 m,
		readTimeout:       timeout,
		replayCache:       replayCache,
		targetIPValidator: onet.RequirePublicIP,
	}
}

// TCPService is a Shadowsocks TCP service that can be started and stopped.
type TCPService interface {
	// SetTargetIPValidator sets the function to be used to validate the target IP addresses.
	SetTargetIPValidator(targetIPValidator onet.TargetIPValidator)
	// Serve adopts the listener, which will be closed before Serve returns.  Serve returns an error unless Stop() was called.
	Serve(listener *net.TCPListener) error
	// Stop closes the listener but does not interfere with existing connections.
	Stop() error
	// GracefulStop calls Stop(), and then blocks until all resources have been cleaned up.
	GracefulStop() error
}

func (s *tcpService) SetTargetIPValidator(targetIPValidator onet.TargetIPValidator) {
	s.targetIPValidator = targetIPValidator
}

// proxyConnection will route the clientConn according to the address read from the connection.
func proxyConnection(clientSSConn onet.DuplexConn, tgtAddr socks.Addr, proxyMetrics *metrics.ProxyMetrics, targetIPValidator onet.TargetIPValidator) *onet.ConnectionError {
	tgtTCPAddr, err := net.ResolveTCPAddr("tcp", tgtAddr.String())
	if err != nil {
		return onet.NewConnectionError("ERR_RESOLVE_ADDRESS", fmt.Sprintf("Failed to resolve target address %v", tgtAddr.String()), err)
	}
	if err := targetIPValidator(tgtTCPAddr.IP); err != nil {
		return err
	}

	tgtTCPConn, err := net.DialTCP("tcp", nil, tgtTCPAddr)
	if err != nil {
		return onet.NewConnectionError("ERR_CONNECT", "Failed to connect to target", err)
	}
	defer tgtTCPConn.Close()
	tgtTCPConn.SetKeepAlive(true)
	tgtConn := metrics.MeasureConn(tgtTCPConn, &proxyMetrics.ProxyTarget, &proxyMetrics.TargetProxy)

	logger.Debugf("proxy %s <-> %s", clientSSConn.RemoteAddr().String(), tgtConn.RemoteAddr().String())
	fromClientErrCh := make(chan error)
	go func() {
		_, fromClientErr := clientSSConn.(io.WriterTo).WriteTo(tgtConn)
		if fromClientErr != nil {
			// Drain to prevent a close on cipher error.
			// TODO: Need to drain the underlying connections instead.
			clientSSConn.(io.WriterTo).WriteTo(ioutil.Discard)
		}
		clientSSConn.CloseRead()
		// Send FIN to target.
		// We must do this after the drain is completed, otherwise the target will close its
		// connection with the proxy, which will, in turn, close the connection with the client.
		tgtConn.CloseWrite()
		fromClientErrCh <- fromClientErr
	}()
	_, fromTargetErr := clientSSConn.(io.ReaderFrom).ReadFrom(tgtConn)
	// Send FIN to client.
	clientSSConn.CloseWrite()
	tgtConn.CloseRead()

	fromClientErr := <-fromClientErrCh
	if fromClientErr != nil {
		return onet.NewConnectionError("ERR_RELAY_CLIENT", "Failed to relay traffic from client", fromClientErr)
	}
	if fromTargetErr != nil {
		return onet.NewConnectionError("ERR_RELAY_TARGET", "Failed to relay traffic from target", fromTargetErr)
	}
	return nil
}

func (s *tcpService) Serve(listener *net.TCPListener) error {
	s.mu.Lock()
	if s.listener != nil {
		s.mu.Unlock()
		listener.Close()
		return errors.New("Serve can only be called once")
	}
	if s.stopped {
		s.mu.Unlock()
		return listener.Close()
	}
	s.listener = listener
	s.running.Add(1)
	s.mu.Unlock()

	defer s.running.Done()
	for {
		clientConn, err := listener.AcceptTCP()
		if err != nil {
			s.mu.RLock()
			stopped := s.stopped
			s.mu.RUnlock()
			if stopped {
				return nil
			}
			logger.Errorf("Accept failed: %v", err)
			continue
		}

		s.running.Add(1)
		go func() {
			defer s.running.Done()
			defer func() {
				if r := recover(); r != nil {
					logger.Errorf("Panic in TCP handler: %v", r)
				}
			}()
			s.handleConnection(listener.Addr().(*net.TCPAddr).Port, clientConn)
		}()
	}
}

func (s *tcpService) handleConnection(listenerPort int, clientTCPConn *net.TCPConn) {
	clientLocation, err := s.m.GetLocation(clientTCPConn.RemoteAddr())
	if err != nil {
		logger.Warningf("Failed location lookup: %v", err)
	}
	logger.Debugf("Got location \"%v\" for IP %v", clientLocation, clientTCPConn.RemoteAddr().String())
	s.m.AddOpenTCPConnection(clientLocation)

	connStart := time.Now()
	clientTCPConn.SetKeepAlive(true)
	// Set a deadline to receive the address to the target.
	clientTCPConn.SetReadDeadline(connStart.Add(s.readTimeout))
	var proxyMetrics metrics.ProxyMetrics
	clientConn := metrics.MeasureConn(clientTCPConn, &proxyMetrics.ProxyClient, &proxyMetrics.ClientProxy)
	cipherEntry, clientReader, clientSalt, timeToCipher, keyErr := findAccessKey(clientConn, remoteIP(clientConn), s.ciphers)

	connError := func() *onet.ConnectionError {
		if keyErr != nil {
			logger.Debugf("Failed to find a valid cipher after reading %v bytes: %v", proxyMetrics.ClientProxy, keyErr)
			const status = "ERR_CIPHER"
			s.absorbProbe(listenerPort, clientConn, clientLocation, status, &proxyMetrics)
			return onet.NewConnectionError(status, "Failed to find a valid cipher", keyErr)
		}

		isServerSalt := cipherEntry.SaltGenerator.IsServerSalt(clientSalt)
		// Only check the cache if findAccessKey succeeded and the salt is unrecognized.
		if isServerSalt || !s.replayCache.Add(cipherEntry.ID, clientSalt) {
			var status string
			if isServerSalt {
				status = "ERR_REPLAY_SERVER"
			} else {
				status = "ERR_REPLAY_CLIENT"
			}
			s.absorbProbe(listenerPort, clientConn, clientLocation, status, &proxyMetrics)
			logger.Debugf(status+": %v in %s sent %d bytes", clientConn.RemoteAddr(), clientLocation, proxyMetrics.ClientProxy)
			return onet.NewConnectionError(status, "Replay detected", nil)
		}

		ssr := ss.NewShadowsocksReader(clientReader, cipherEntry.Cipher)
		tgtAddr, err := socks.ReadAddr(ssr)
		if err != nil {
			// Drain to prevent a close on cipher error.
			clientConn.(io.WriterTo).WriteTo(ioutil.Discard)
			return onet.NewConnectionError("ERR_READ_ADDRESS", "Failed to get target address", err)
		}

		// Clear the deadline for the target address
		clientConn.SetReadDeadline(time.Time{})
		ssw := ss.NewShadowsocksWriter(clientConn, cipherEntry.Cipher)
		ssw.SetSaltGenerator(cipherEntry.SaltGenerator)
		clientSSConn := onet.WrapConn(clientConn, ssr, ssw)
		return proxyConnection(clientSSConn, tgtAddr, &proxyMetrics, s.targetIPValidator)
	}()

	connDuration := time.Now().Sub(connStart)
	status := "OK"
	if connError != nil {
		logger.Debugf("TCP Error: %v: %v", connError.Message, connError.Cause)
		status = connError.Status
	}
	var id string
	if cipherEntry != nil {
		id = cipherEntry.ID
	}
	s.m.AddClosedTCPConnection(clientLocation, id, status, proxyMetrics, timeToCipher, connDuration)
	clientConn.Close() // Closing after the metrics are added aids integration testing.
	logger.Debugf("Done with status %v, duration %v", status, connDuration)
}

// Keep the connection open until we hit the authentication deadline to protect against probing attacks
// `proxyMetrics` is a pointer because its value is being mutated by `clientConn`.
func (s *tcpService) absorbProbe(listenerPort int, clientConn io.ReadCloser, clientLocation, status string, proxyMetrics *metrics.ProxyMetrics) {
	_, drainErr := io.Copy(ioutil.Discard, clientConn) // drain socket
	drainResult := drainErrToString(drainErr)
	logger.Debugf("Drain error: %v, drain result: %v", drainErr, drainResult)
	s.m.AddTCPProbe(clientLocation, status, drainResult, listenerPort, *proxyMetrics)
}

func drainErrToString(drainErr error) string {
	netErr, ok := drainErr.(net.Error)
	switch {
	case drainErr == nil:
		return "eof"
	case ok && netErr.Timeout():
		return "timeout"
	default:
		return "other"
	}
}

func (s *tcpService) Stop() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.stopped = true
	if s.listener == nil {
		return nil
	}
	return s.listener.Close()
}

func (s *tcpService) GracefulStop() error {
	err := s.Stop()
	s.running.Wait()
	return err
}
<|endoftext|>


<commit_before>//
// SPDX-License-Identifier: BSD-3-Clause
//

package gofish

import (
	"encoding/json"

	"github.com/stmcginnis/gofish/common"
	"github.com/stmcginnis/gofish/redfish"
	"github.com/stmcginnis/gofish/swordfish"
)

// Expand shall contain information about the support of the $expand query
// parameter by the service.
type Expand struct {
	// ExpandAll shall be a boolean indicating whether this service supports the
	// use of asterisk (expand all entries) as a value for the $expand query
	// parameter as described by the specification.
	ExpandAll bool
	// Levels shall be a boolean indicating whether this service supports the
	// use of $levels as a value for the $expand query parameter as described by
	// the specification.
	Levels bool
	// Links shall be a boolean indicating whether this service supports the use
	// of tilde (expand only entries in the Links section) as a value for the
	// $expand query parameter as described by the specification.
	Links bool
	// MaxLevels shall be the maximum value of the $levels qualifier supported
	// by the service and shall only be included if the value of the Levels
	// property is true.
	MaxLevels int
	// NoLinks shall be a boolean indicating whether this service supports the
	// use of period (expand only entries not in the Links section) as a value
	// for the $expand query parameter as described by the specification.
	NoLinks bool
}

// ProtocolFeaturesSupported contains information about protocol features
// supported by the service.
type ProtocolFeaturesSupported struct {
	// ExcerptQuery shall be a boolean indicating whether this service supports
	// the use of the 'excerpt' query parameter as described by the
	// specification.
	ExcerptQuery bool
	// ExpandQuery shall contain information about the support of the $expand
	// query parameter by the service.
	ExpandQuery Expand
	// FilterQuery shall be a boolean indicating whether this service supports
	// the use of the $filter query parameter as described by the specification.
	FilterQuery bool
	// OnlyMemberQuery shall be a boolean indicating whether this service
	// supports the use of the 'only' query parameter as described by the
	// specification.
	OnlyMemberQuery bool
	// SelectQuery shall be a boolean indicating whether this service supports
	// the use of the $select query parameter as described by the specification.
	SelectQuery bool
}

// Service represents the root Redfish service. All values for resources
// described by this schema shall comply to the requirements as described in the
// Redfish specification.
type Service struct {
	common.Entity

	// ODataContext is the odata context.
	ODataContext string `json:"@odata.context"`
	// ODataID is the odata identifier.
	ODataID string `json:"@odata.id"`
	// ODataType is the odata type.
	ODataType string `json:"@odata.type"`
	// AccountService shall only contain a reference to a resource that complies
	// to the AccountService schema.
	accountService string
	// CertificateService shall be a link to the CertificateService.
	certificateService string
	// Chassis shall only contain a reference to a collection of resources that
	// comply to the Chassis schema.
	chassis string
	// CompositionService shall only contain a reference to a resource that
	// complies to the CompositionService schema.
	compositionService string
	// Description provides a description of this resource.
	Description string
	// EventService shall only contain a reference to a resource that complies
	// to the EventService schema.
	eventService string
	// Fabrics shall contain references to all Fabric instances.
	fabrics string
	// JobService shall only contain a reference to a resource that conforms to
	// the JobService schema.
	jobService string
	// JsonSchemas shall only contain a reference to a collection of resources
	// that comply to the SchemaFile schema where the files are Json-Schema
	// files.
	jsonSchemas string
	// Managers shall only contain a reference to a collection of resources that
	// comply to the Managers schema.
	managers string
	// Product shall include the name of the product represented by this Redfish
	// service.
	Product string
	// ProtocolFeaturesSupported contains information about protocol features
	// supported by the service.
	ProtocolFeaturesSupported ProtocolFeaturesSupported
	// RedfishVersion shall represent the version of the Redfish service. The
	// format of this string shall be of the format
	// majorversion.minorversion.errata in compliance with Protocol Version
	// section of the Redfish specification.
	RedfishVersion string
	// Registries shall contain a reference to Message Registry.
	registries string
	// ResourceBlocks shall contain references to all Resource Block instances.
	resourceBlocks string
	// SessionService shall only contain a reference to a resource that complies
	// to the SessionService schema.
	sessionService string
	// StorageServices shall contain references to all StorageService instances.
	storageServices string
	// StorageSystems shall contain computer systems that act as storage
	// servers. The HostingRoles attribute of each such computer system shall
	// have an entry for StorageServer.
	storageSystems string
	// Systems shall only contain a reference to a collection of resources that
	// comply to the Systems schema.
	systems string
	// Tasks shall only contain a reference to a resource that complies to the
	// TaskService schema.
	tasks string
	// TelemetryService shall be a link to the TelemetryService.
	telemetryService string
	// UUID shall be an exact match of the UUID value returned in a 200OK from
	// an SSDP M-SEARCH request during discovery. RFC4122 describes methods that
	// can be used to create a UUID value. The value should be considered to be
	// opaque. Client software should only treat the overall value as a
	// universally unique identifier and should not interpret any sub-fields
	// within the UUID.
	UUID string
	// UpdateService shall only contain a reference to a resource that complies
	// to the UpdateService schema.
	updateService string
	// Vendor shall include the name of the manufacturer or vendor represented
	// by this Redfish service. If this property is supported, the vendor name
	// shall not be included in the value of the Product property.
	Vendor string
	// Sessions shall contain the link to a collection of Sessions.
	sessions string
}

// UnmarshalJSON unmarshals a Service object from the raw JSON.
func (serviceroot *Service) UnmarshalJSON(b []byte) error {
	type temp Service
	var t struct {
		temp
		CertificateService common.Link
		Chassis            common.Link
		Managers           common.Link
		Tasks              common.Link
		StorageServices    common.Link
		StorageSystems     common.Link
		AccountService     common.Link
		EventService       common.Link
		Registries         common.Link
		Systems            common.Link
		CompositionService common.Link
		Fabrics            common.Link
		JobService         common.Link
		JSONSchemas        common.Link `json:"JsonSchemas"`
		ResourceBlocks     common.Link
		SessionService     common.Link
		TelemetryService   common.Link
		UpdateService      common.Link
		Links              struct {
			Sessions common.Link
		}
	}

	err := json.Unmarshal(b, &t)
	if err != nil {
		return err
	}

	// Extract the links to other entities for later
	*serviceroot = Service(t.temp)
	serviceroot.certificateService = string(t.CertificateService)
	serviceroot.chassis = string(t.Chassis)
	serviceroot.managers = string(t.Managers)
	serviceroot.tasks = string(t.Tasks)
	serviceroot.sessions = string(t.Links.Sessions)
	serviceroot.storageServices = string(t.StorageServices)
	serviceroot.storageSystems = string(t.StorageSystems)
	serviceroot.accountService = string(t.AccountService)
	serviceroot.eventService = string(t.EventService)
	serviceroot.registries = string(t.Registries)
	serviceroot.systems = string(t.Systems)
	serviceroot.compositionService = string(t.CompositionService)
	serviceroot.fabrics = string(t.Fabrics)
	serviceroot.jobService = string(t.JobService)
	serviceroot.jsonSchemas = string(t.JSONSchemas)
	serviceroot.resourceBlocks = string(t.ResourceBlocks)
	serviceroot.sessionService = string(t.SessionService)
	serviceroot.telemetryService = string(t.TelemetryService)
	serviceroot.updateService = string(t.UpdateService)

	return nil
}

// ServiceRoot will get a Service instance from the service.
func ServiceRoot(c common.Client) (*Service, error) {
	resp, err := c.Get(common.DefaultServiceRoot)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	var serviceroot Service
	err = json.NewDecoder(resp.Body).Decode(&serviceroot)
	if err != nil {
		return nil, err
	}

	serviceroot.SetClient(c)
	return &serviceroot, nil
}

// Chassis gets the chassis instances managed by this service.
func (serviceroot *Service) Chassis() ([]*redfish.Chassis, error) {
	return redfish.ListReferencedChassis(serviceroot.Client, serviceroot.chassis)
}

// Managers gets the manager instances of this service.
func (serviceroot *Service) Managers() ([]*redfish.Manager, error) {
	return redfish.ListReferencedManagers(serviceroot.Client, serviceroot.managers)
}

// StorageSystems gets the storage system instances managed by this service.
func (serviceroot *Service) StorageSystems() ([]*swordfish.StorageSystem, error) {
	return swordfish.ListReferencedStorageSystems(serviceroot.Client, serviceroot.storageSystems)
}

// StorageServices gets the Swordfish storage services
func (serviceroot *Service) StorageServices() ([]*swordfish.StorageService, error) {
	return swordfish.ListReferencedStorageServices(serviceroot.Client, serviceroot.storageServices)
}

// Tasks gets the system's tasks
func (serviceroot *Service) Tasks() ([]*redfish.Task, error) {
	return redfish.ListReferencedTasks(serviceroot.Client, serviceroot.tasks)
}

// CreateSession creates a new session and returns the token and id
func (serviceroot *Service) CreateSession(username string, password string) (*redfish.AuthToken, error) {
	return redfish.CreateSession(serviceroot.Client, serviceroot.sessions, username, password)
}

// Sessions gets the system's active sessions
func (serviceroot *Service) Sessions() ([]*redfish.Session, error) {
	return redfish.ListReferencedSessions(serviceroot.Client, serviceroot.sessions)
}

// DeleteSession logout the specified session
func (serviceroot *Service) DeleteSession(url string) error {
	return redfish.DeleteSession(serviceroot.Client, url)
}

// AccountService gets the Redfish AccountService
func (serviceroot *Service) AccountService() (*redfish.AccountService, error) {
	return redfish.GetAccountService(serviceroot.Client, serviceroot.accountService)
}

// EventService gets the Redfish EventService
func (serviceroot *Service) EventService() (*redfish.EventService, error) {
	return redfish.GetEventService(serviceroot.Client, serviceroot.eventService)
}

// Registries gets the Redfish Registries
func (serviceroot *Service) Registries() ([]*redfish.MessageRegistryFile, error) {
	return redfish.ListReferencedMessageRegistryFiles(serviceroot.Client, serviceroot.registries)
}

// MessageRegistries gets all the available message registries in all languages
func (serviceroot *Service) MessageRegistries() ([]*redfish.MessageRegistry, error) {
	return redfish.ListReferencedMessageRegistries(serviceroot.Client, serviceroot.registries)
}

// MessageRegistriesByLanguage gets the message registries by language.
// language is the RFC5646-conformant language code for the message registry, for example: "en".
func (serviceroot *Service) MessageRegistriesByLanguage(language string) ([]*redfish.MessageRegistry, error) {
	return redfish.ListReferencedMessageRegistriesByLanguage(serviceroot.Client, serviceroot.registries, language)
}

// MessageRegistryByLanguage gets a specific message registry by language.
// registry is used to identify the correct Message Registry file and it shall
// contain the Message Registry name and it major and minor versions, as defined
// by the Redfish Specification, for example: "Alert.1.0.0".
// language is the RFC5646-conformant language code for the message registry, for example: "en".
func (serviceroot *Service) MessageRegistryByLanguage(registry string, language string) (*redfish.MessageRegistry, error) {
	return redfish.GetMessageRegistryByLanguage(serviceroot.Client, serviceroot.registries, registry, language)
}

// MessageByLanguage tries to find and get the message in the correct language from the informed messageID.
// messageID is the key used to find the registry, version and message, for example: "Alert.1.0.LanDisconnect"
//  - The segment before the 1st period is the Registry Name (Registry Prefix): Alert
//  - The segment between the 1st and 2nd period is the major version: 1
//  - The segment between the 2nd and 3rd period is the minor version: 0
//  - The segment after the 3rd period is the Message Identifier in the Registry: LanDisconnect
// language is the RFC5646-conformant language code for the message registry, for example: "en".
func (serviceroot *Service) MessageByLanguage(messageID string, language string) (*redfish.MessageRegistryMessage, error) {
	return redfish.GetMessageFromMessageRegistryByLanguage(serviceroot.Client, serviceroot.registries, messageID, language)
}

// Systems get the system instances from the service
func (serviceroot *Service) Systems() ([]*redfish.ComputerSystem, error) {
	return redfish.ListReferencedComputerSystems(serviceroot.Client, serviceroot.systems)
}

// CompositionService gets the composition service instance
func (serviceroot *Service) CompositionService() (*redfish.CompositionService, error) {
	return redfish.GetCompositionService(serviceroot.Client, serviceroot.compositionService)
}

// UpdateService gets the update service instance
func (serviceroot *Service) UpdateService() (*redfish.UpdateService, error) {
	return redfish.GetUpdateService(serviceroot.Client, serviceroot.updateService)
}
<commit_msg>The search for the Redfish message registry by registry/language can be very slow depending on the environment (existing number of registries / size of the registries). Created new function to get the MessageRegistry directly from the uri avoiding doing unecessary searches to find it.<commit_after>//
// SPDX-License-Identifier: BSD-3-Clause
//

package gofish

import (
	"encoding/json"

	"github.com/stmcginnis/gofish/common"
	"github.com/stmcginnis/gofish/redfish"
	"github.com/stmcginnis/gofish/swordfish"
)

// Expand shall contain information about the support of the $expand query
// parameter by the service.
type Expand struct {
	// ExpandAll shall be a boolean indicating whether this service supports the
	// use of asterisk (expand all entries) as a value for the $expand query
	// parameter as described by the specification.
	ExpandAll bool
	// Levels shall be a boolean indicating whether this service supports the
	// use of $levels as a value for the $expand query parameter as described by
	// the specification.
	Levels bool
	// Links shall be a boolean indicating whether this service supports the use
	// of tilde (expand only entries in the Links section) as a value for the
	// $expand query parameter as described by the specification.
	Links bool
	// MaxLevels shall be the maximum value of the $levels qualifier supported
	// by the service and shall only be included if the value of the Levels
	// property is true.
	MaxLevels int
	// NoLinks shall be a boolean indicating whether this service supports the
	// use of period (expand only entries not in the Links section) as a value
	// for the $expand query parameter as described by the specification.
	NoLinks bool
}

// ProtocolFeaturesSupported contains information about protocol features
// supported by the service.
type ProtocolFeaturesSupported struct {
	// ExcerptQuery shall be a boolean indicating whether this service supports
	// the use of the 'excerpt' query parameter as described by the
	// specification.
	ExcerptQuery bool
	// ExpandQuery shall contain information about the support of the $expand
	// query parameter by the service.
	ExpandQuery Expand
	// FilterQuery shall be a boolean indicating whether this service supports
	// the use of the $filter query parameter as described by the specification.
	FilterQuery bool
	// OnlyMemberQuery shall be a boolean indicating whether this service
	// supports the use of the 'only' query parameter as described by the
	// specification.
	OnlyMemberQuery bool
	// SelectQuery shall be a boolean indicating whether this service supports
	// the use of the $select query parameter as described by the specification.
	SelectQuery bool
}

// Service represents the root Redfish service. All values for resources
// described by this schema shall comply to the requirements as described in the
// Redfish specification.
type Service struct {
	common.Entity

	// ODataContext is the odata context.
	ODataContext string `json:"@odata.context"`
	// ODataID is the odata identifier.
	ODataID string `json:"@odata.id"`
	// ODataType is the odata type.
	ODataType string `json:"@odata.type"`
	// AccountService shall only contain a reference to a resource that complies
	// to the AccountService schema.
	accountService string
	// CertificateService shall be a link to the CertificateService.
	certificateService string
	// Chassis shall only contain a reference to a collection of resources that
	// comply to the Chassis schema.
	chassis string
	// CompositionService shall only contain a reference to a resource that
	// complies to the CompositionService schema.
	compositionService string
	// Description provides a description of this resource.
	Description string
	// EventService shall only contain a reference to a resource that complies
	// to the EventService schema.
	eventService string
	// Fabrics shall contain references to all Fabric instances.
	fabrics string
	// JobService shall only contain a reference to a resource that conforms to
	// the JobService schema.
	jobService string
	// JsonSchemas shall only contain a reference to a collection of resources
	// that comply to the SchemaFile schema where the files are Json-Schema
	// files.
	jsonSchemas string
	// Managers shall only contain a reference to a collection of resources that
	// comply to the Managers schema.
	managers string
	// Product shall include the name of the product represented by this Redfish
	// service.
	Product string
	// ProtocolFeaturesSupported contains information about protocol features
	// supported by the service.
	ProtocolFeaturesSupported ProtocolFeaturesSupported
	// RedfishVersion shall represent the version of the Redfish service. The
	// format of this string shall be of the format
	// majorversion.minorversion.errata in compliance with Protocol Version
	// section of the Redfish specification.
	RedfishVersion string
	// Registries shall contain a reference to Message Registry.
	registries string
	// ResourceBlocks shall contain references to all Resource Block instances.
	resourceBlocks string
	// SessionService shall only contain a reference to a resource that complies
	// to the SessionService schema.
	sessionService string
	// StorageServices shall contain references to all StorageService instances.
	storageServices string
	// StorageSystems shall contain computer systems that act as storage
	// servers. The HostingRoles attribute of each such computer system shall
	// have an entry for StorageServer.
	storageSystems string
	// Systems shall only contain a reference to a collection of resources that
	// comply to the Systems schema.
	systems string
	// Tasks shall only contain a reference to a resource that complies to the
	// TaskService schema.
	tasks string
	// TelemetryService shall be a link to the TelemetryService.
	telemetryService string
	// UUID shall be an exact match of the UUID value returned in a 200OK from
	// an SSDP M-SEARCH request during discovery. RFC4122 describes methods that
	// can be used to create a UUID value. The value should be considered to be
	// opaque. Client software should only treat the overall value as a
	// universally unique identifier and should not interpret any sub-fields
	// within the UUID.
	UUID string
	// UpdateService shall only contain a reference to a resource that complies
	// to the UpdateService schema.
	updateService string
	// Vendor shall include the name of the manufacturer or vendor represented
	// by this Redfish service. If this property is supported, the vendor name
	// shall not be included in the value of the Product property.
	Vendor string
	// Sessions shall contain the link to a collection of Sessions.
	sessions string
}

// UnmarshalJSON unmarshals a Service object from the raw JSON.
func (serviceroot *Service) UnmarshalJSON(b []byte) error {
	type temp Service
	var t struct {
		temp
		CertificateService common.Link
		Chassis            common.Link
		Managers           common.Link
		Tasks              common.Link
		StorageServices    common.Link
		StorageSystems     common.Link
		AccountService     common.Link
		EventService       common.Link
		Registries         common.Link
		Systems            common.Link
		CompositionService common.Link
		Fabrics            common.Link
		JobService         common.Link
		JSONSchemas        common.Link `json:"JsonSchemas"`
		ResourceBlocks     common.Link
		SessionService     common.Link
		TelemetryService   common.Link
		UpdateService      common.Link
		Links              struct {
			Sessions common.Link
		}
	}

	err := json.Unmarshal(b, &t)
	if err != nil {
		return err
	}

	// Extract the links to other entities for later
	*serviceroot = Service(t.temp)
	serviceroot.certificateService = string(t.CertificateService)
	serviceroot.chassis = string(t.Chassis)
	serviceroot.managers = string(t.Managers)
	serviceroot.tasks = string(t.Tasks)
	serviceroot.sessions = string(t.Links.Sessions)
	serviceroot.storageServices = string(t.StorageServices)
	serviceroot.storageSystems = string(t.StorageSystems)
	serviceroot.accountService = string(t.AccountService)
	serviceroot.eventService = string(t.EventService)
	serviceroot.registries = string(t.Registries)
	serviceroot.systems = string(t.Systems)
	serviceroot.compositionService = string(t.CompositionService)
	serviceroot.fabrics = string(t.Fabrics)
	serviceroot.jobService = string(t.JobService)
	serviceroot.jsonSchemas = string(t.JSONSchemas)
	serviceroot.resourceBlocks = string(t.ResourceBlocks)
	serviceroot.sessionService = string(t.SessionService)
	serviceroot.telemetryService = string(t.TelemetryService)
	serviceroot.updateService = string(t.UpdateService)

	return nil
}

// ServiceRoot will get a Service instance from the service.
func ServiceRoot(c common.Client) (*Service, error) {
	resp, err := c.Get(common.DefaultServiceRoot)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	var serviceroot Service
	err = json.NewDecoder(resp.Body).Decode(&serviceroot)
	if err != nil {
		return nil, err
	}

	serviceroot.SetClient(c)
	return &serviceroot, nil
}

// Chassis gets the chassis instances managed by this service.
func (serviceroot *Service) Chassis() ([]*redfish.Chassis, error) {
	return redfish.ListReferencedChassis(serviceroot.Client, serviceroot.chassis)
}

// Managers gets the manager instances of this service.
func (serviceroot *Service) Managers() ([]*redfish.Manager, error) {
	return redfish.ListReferencedManagers(serviceroot.Client, serviceroot.managers)
}

// StorageSystems gets the storage system instances managed by this service.
func (serviceroot *Service) StorageSystems() ([]*swordfish.StorageSystem, error) {
	return swordfish.ListReferencedStorageSystems(serviceroot.Client, serviceroot.storageSystems)
}

// StorageServices gets the Swordfish storage services
func (serviceroot *Service) StorageServices() ([]*swordfish.StorageService, error) {
	return swordfish.ListReferencedStorageServices(serviceroot.Client, serviceroot.storageServices)
}

// Tasks gets the system's tasks
func (serviceroot *Service) Tasks() ([]*redfish.Task, error) {
	return redfish.ListReferencedTasks(serviceroot.Client, serviceroot.tasks)
}

// CreateSession creates a new session and returns the token and id
func (serviceroot *Service) CreateSession(username string, password string) (*redfish.AuthToken, error) {
	return redfish.CreateSession(serviceroot.Client, serviceroot.sessions, username, password)
}

// Sessions gets the system's active sessions
func (serviceroot *Service) Sessions() ([]*redfish.Session, error) {
	return redfish.ListReferencedSessions(serviceroot.Client, serviceroot.sessions)
}

// DeleteSession logout the specified session
func (serviceroot *Service) DeleteSession(url string) error {
	return redfish.DeleteSession(serviceroot.Client, url)
}

// AccountService gets the Redfish AccountService
func (serviceroot *Service) AccountService() (*redfish.AccountService, error) {
	return redfish.GetAccountService(serviceroot.Client, serviceroot.accountService)
}

// EventService gets the Redfish EventService
func (serviceroot *Service) EventService() (*redfish.EventService, error) {
	return redfish.GetEventService(serviceroot.Client, serviceroot.eventService)
}

// Registries gets the Redfish Registries
func (serviceroot *Service) Registries() ([]*redfish.MessageRegistryFile, error) {
	return redfish.ListReferencedMessageRegistryFiles(serviceroot.Client, serviceroot.registries)
}

// MessageRegistries gets all the available message registries in all languages
func (serviceroot *Service) MessageRegistries() ([]*redfish.MessageRegistry, error) {
	return redfish.ListReferencedMessageRegistries(serviceroot.Client, serviceroot.registries)
}

// MessageRegistry gets a specific message registry.
// uri is the uri for the message registry
func (serviceroot *Service) MessageRegistry(uri string) (*redfish.MessageRegistry, error) {
	return redfish.GetMessageRegistry(serviceroot.Client, uri)
}

// MessageRegistriesByLanguage gets the message registries by language.
// language is the RFC5646-conformant language code for the message registry, for example: "en".
func (serviceroot *Service) MessageRegistriesByLanguage(language string) ([]*redfish.MessageRegistry, error) {
	return redfish.ListReferencedMessageRegistriesByLanguage(serviceroot.Client, serviceroot.registries, language)
}

// MessageRegistryByLanguage gets a specific message registry by language.
// registry is used to identify the correct Message Registry file and it shall
// contain the Message Registry name and it major and minor versions, as defined
// by the Redfish Specification, for example: "Alert.1.0.0".
// language is the RFC5646-conformant language code for the message registry, for example: "en".
func (serviceroot *Service) MessageRegistryByLanguage(registry string, language string) (*redfish.MessageRegistry, error) {
	return redfish.GetMessageRegistryByLanguage(serviceroot.Client, serviceroot.registries, registry, language)
}

// MessageByLanguage tries to find and get the message in the correct language from the informed messageID.
// messageID is the key used to find the registry, version and message, for example: "Alert.1.0.LanDisconnect"
//  - The segment before the 1st period is the Registry Name (Registry Prefix): Alert
//  - The segment between the 1st and 2nd period is the major version: 1
//  - The segment between the 2nd and 3rd period is the minor version: 0
//  - The segment after the 3rd period is the Message Identifier in the Registry: LanDisconnect
// language is the RFC5646-conformant language code for the message registry, for example: "en".
func (serviceroot *Service) MessageByLanguage(messageID string, language string) (*redfish.MessageRegistryMessage, error) {
	return redfish.GetMessageFromMessageRegistryByLanguage(serviceroot.Client, serviceroot.registries, messageID, language)
}

// Systems get the system instances from the service
func (serviceroot *Service) Systems() ([]*redfish.ComputerSystem, error) {
	return redfish.ListReferencedComputerSystems(serviceroot.Client, serviceroot.systems)
}

// CompositionService gets the composition service instance
func (serviceroot *Service) CompositionService() (*redfish.CompositionService, error) {
	return redfish.GetCompositionService(serviceroot.Client, serviceroot.compositionService)
}

// UpdateService gets the update service instance
func (serviceroot *Service) UpdateService() (*redfish.UpdateService, error) {
	return redfish.GetUpdateService(serviceroot.Client, serviceroot.updateService)
}
<|endoftext|>


<commit_before>package simplebp

import (
	"github.com/google/blueprint"
	"github.com/google/blueprint/pathtools"

	"bytes"
	"os"
	"path/filepath"
	"strings"
	"text/template"
)

var (
	scriptRule = pctx.StaticRule("script",
		blueprint.RuleParams{
			Command:     "$script $args",
			Description: "RUN $script",
		},
		"script", "args")
)

type ScriptModule struct {
	properties struct {
		Script string
		Inputs []string
		Output string
		Args   string
	}
}

func NewScript() (blueprint.Module, []interface{}) {
	module := new(ScriptModule)
	properties := &module.properties
	return module, []interface{}{properties}
}

type scriptInput struct {
	Name      string
	Basename  string
	Extension string
}

type scriptArgs struct {
	Input  scriptInput
	Output string
}

func (m *ScriptModule) GenerateBuildActions(ctx blueprint.ModuleContext) {
	config := ctx.Config().(*config)

	var scriptPath string
	if s := m.properties.Script; strings.HasPrefix(s, "//") {
		scriptPath = filepath.Join(config.srcDir, s[2:])
	} else {
		scriptPath = filepath.Join(ctx.ModuleDir(), s)
	}

	if stat, err := os.Stat(scriptPath); err != nil {
		ctx.ModuleErrorf("Could not stat %v: %v", scriptPath, err)
		return
	} else if stat.Mode()&0111 == 0 {
		ctx.ModuleErrorf("%s is not an executable", scriptPath)
		return
	}

	srcs := pathtools.PrefixPaths(m.properties.Inputs, ctx.ModuleDir())

	argsTmpl, argsErr := template.New("args").Parse(m.properties.Args)
	if argsErr != nil {
		ctx.ModuleErrorf("Could not parse script args: %v", argsErr)
		return
	}
	outTmpl, outErr := template.New("out").Parse(m.properties.Output)
	if outErr != nil {
		ctx.ModuleErrorf("Could not parse output template: %v", outErr)
		return
	}

	for _, s := range srcs {
		args := &scriptArgs{
			Input: scriptInput{
				Name:      s,
				Basename:  strings.TrimSuffix(s, filepath.Ext(s)),
				Extension: filepath.Ext(s),
			},
		}

		outBuf := &bytes.Buffer{}
		if err := outTmpl.Execute(outBuf, args); err != nil {
			ctx.ModuleErrorf("Could not generate output: %v", err)
			return
		}

		args.Output = filepath.Join(config.buildDir, outBuf.String())
		argsBuf := &bytes.Buffer{}
		if err := argsTmpl.Execute(argsBuf, args); err != nil {
			ctx.ModuleErrorf("Could not generate args: %v", err)
			return
		}

		ctx.Build(pctx, blueprint.BuildParams{
			Rule:    scriptRule,
			Inputs:  []string{s},
			Outputs: []string{args.Output},
                        Implicits: []string{scriptPath},
			Args: map[string]string{
				"script": scriptPath,
				"args":   argsBuf.String(),
			},
		})
	}
}
<commit_msg>run gofmt<commit_after>package simplebp

import (
	"github.com/google/blueprint"
	"github.com/google/blueprint/pathtools"

	"bytes"
	"os"
	"path/filepath"
	"strings"
	"text/template"
)

var (
	scriptRule = pctx.StaticRule("script",
		blueprint.RuleParams{
			Command:     "$script $args",
			Description: "RUN $script",
		},
		"script", "args")
)

type ScriptModule struct {
	properties struct {
		Script string
		Inputs []string
		Output string
		Args   string
	}
}

func NewScript() (blueprint.Module, []interface{}) {
	module := new(ScriptModule)
	properties := &module.properties
	return module, []interface{}{properties}
}

type scriptInput struct {
	Name      string
	Basename  string
	Extension string
}

type scriptArgs struct {
	Input  scriptInput
	Output string
}

func (m *ScriptModule) GenerateBuildActions(ctx blueprint.ModuleContext) {
	config := ctx.Config().(*config)

	var scriptPath string
	if s := m.properties.Script; strings.HasPrefix(s, "//") {
		scriptPath = filepath.Join(config.srcDir, s[2:])
	} else {
		scriptPath = filepath.Join(ctx.ModuleDir(), s)
	}

	if stat, err := os.Stat(scriptPath); err != nil {
		ctx.ModuleErrorf("Could not stat %v: %v", scriptPath, err)
		return
	} else if stat.Mode()&0111 == 0 {
		ctx.ModuleErrorf("%s is not an executable", scriptPath)
		return
	}

	srcs := pathtools.PrefixPaths(m.properties.Inputs, ctx.ModuleDir())

	argsTmpl, argsErr := template.New("args").Parse(m.properties.Args)
	if argsErr != nil {
		ctx.ModuleErrorf("Could not parse script args: %v", argsErr)
		return
	}
	outTmpl, outErr := template.New("out").Parse(m.properties.Output)
	if outErr != nil {
		ctx.ModuleErrorf("Could not parse output template: %v", outErr)
		return
	}

	for _, s := range srcs {
		args := &scriptArgs{
			Input: scriptInput{
				Name:      s,
				Basename:  strings.TrimSuffix(s, filepath.Ext(s)),
				Extension: filepath.Ext(s),
			},
		}

		outBuf := &bytes.Buffer{}
		if err := outTmpl.Execute(outBuf, args); err != nil {
			ctx.ModuleErrorf("Could not generate output: %v", err)
			return
		}

		args.Output = filepath.Join(config.buildDir, outBuf.String())
		argsBuf := &bytes.Buffer{}
		if err := argsTmpl.Execute(argsBuf, args); err != nil {
			ctx.ModuleErrorf("Could not generate args: %v", err)
			return
		}

		ctx.Build(pctx, blueprint.BuildParams{
			Rule:      scriptRule,
			Inputs:    []string{s},
			Outputs:   []string{args.Output},
			Implicits: []string{scriptPath},
			Args: map[string]string{
				"script": scriptPath,
				"args":   argsBuf.String(),
			},
		})
	}
}
<|endoftext|>


<commit_before>package images

import(
  "io"
  "io/ioutil"
  "encoding/json"
  "github.com/ricallinson/forgery"
  "github.com/spacedock-io/registry/db"
  "github.com/spacedock-io/registry/models"
  "github.com/spacedock-io/registry/cloudfiles"
)

func GetJson(req *f.Request, res *f.Response) {
  var image models.Image
  q := db.DB.Where(&models.Image{Uuid: req.Params["id"]}).First(&image)
  if q.Error != nil {
    res.Send(404)
    return
  }

  res.Set("X-Docker-Size", string(image.Size))
  res.Set("X-Docker-Checksum", image.Checksum)

  res.Send(image.Json)
}

func PutJson(req *f.Request, res *f.Response) {
  var image models.Image
  var err error

  q := db.DB.Where(&models.Image{Uuid: req.Params["id"]}).First(&image)
  if !q.RecordNotFound() && q.Error != nil {
    res.Send(404)
    return
  }

  image.Json, err = ioutil.ReadAll(req.Request.Request.Body)

  if err == nil {
    db.DB.Save(&image)
  } else {
    res.Send(500)
  }
}

func GetLayer(req *f.Request, res *f.Response) {
  _, err := cloudfiles.Cloudfiles.ObjectGet(
    "default", req.Params["id"], res.Response.Writer, true, nil)
  if err == nil {
    res.Send(200)
  } else { res.Send(500) }
}

func PutLayer(req *f.Request, res *f.Response) {
  obj, err := cloudfiles.Cloudfiles.ObjectCreate(
    "default", req.Params["id"], true, "", "", nil)
  if err == nil {
    io.Copy(obj, req.Request.Request.Body)
    res.Send(200)
  } else { res.Send(500) }
}

func GetAncestry(req *f.Request, res *f.Response) {
  var image models.Image
  q := db.DB.First(&models.Image{Uuid: req.Params["id"]}).First(&image)
  if q.Error != nil {
    res.Send(404)
    return
  }

  data, err := json.Marshal(image.Ancestry)

  if err == nil {
    res.Send(data)
  } else { res.Send(500) }
}
<commit_msg>Let's see what the query is returning<commit_after>package images

import(
  "fmt"
  "io"
  "io/ioutil"
  "encoding/json"
  "github.com/ricallinson/forgery"
  "github.com/spacedock-io/registry/db"
  "github.com/spacedock-io/registry/models"
  "github.com/spacedock-io/registry/cloudfiles"
)

func GetJson(req *f.Request, res *f.Response) {
  var image models.Image
  q := db.DB.Where(&models.Image{Uuid: req.Params["id"]}).First(&image)
  if q.Error != nil {
    res.Send(404)
    return
  }

  res.Set("X-Docker-Size", string(image.Size))
  res.Set("X-Docker-Checksum", image.Checksum)

  res.Send(image.Json)
}

func PutJson(req *f.Request, res *f.Response) {
  var image models.Image
  var err error

  q := db.DB.Where(&models.Image{Uuid: req.Params["id"]}).First(&image)
  fmt.Printf("q: %+v\n", q)
  if !q.RecordNotFound() && q.Error != nil {
    res.Send(404)
    return
  }

  image.Json, err = ioutil.ReadAll(req.Request.Request.Body)

  if err == nil {
    db.DB.Save(&image)
  } else {
    res.Send(500)
  }
}

func GetLayer(req *f.Request, res *f.Response) {
  _, err := cloudfiles.Cloudfiles.ObjectGet(
    "default", req.Params["id"], res.Response.Writer, true, nil)
  if err == nil {
    res.Send(200)
  } else { res.Send(500) }
}

func PutLayer(req *f.Request, res *f.Response) {
  obj, err := cloudfiles.Cloudfiles.ObjectCreate(
    "default", req.Params["id"], true, "", "", nil)
  if err == nil {
    io.Copy(obj, req.Request.Request.Body)
    res.Send(200)
  } else { res.Send(500) }
}

func GetAncestry(req *f.Request, res *f.Response) {
  var image models.Image
  q := db.DB.First(&models.Image{Uuid: req.Params["id"]}).First(&image)
  if q.Error != nil {
    res.Send(404)
    return
  }

  data, err := json.Marshal(image.Ancestry)

  if err == nil {
    res.Send(data)
  } else { res.Send(500) }
}
<|endoftext|>


<commit_before>package memory

func createTestCache() *MemoryCache {
	return New(20 * 1024 * 1024)
}
<commit_msg>add tests for MemoryCache<commit_after>package memory

import (
	"github.com/pierrre/imageserver/cache/cachetest"
	"testing"
)

func TestGetSet(t *testing.T) {
	cache := createTestCache()

	cachetest.CacheTestGetSetAllImages(t, cache)
}

func TestGetErrorMiss(t *testing.T) {
	cache := createTestCache()

	cachetest.CacheTestGetErrorMiss(t, cache)
}

func createTestCache() *MemoryCache {
	return New(20 * 1024 * 1024)
}
<|endoftext|>


<commit_before>package parse

import (
	"strings"
	"testing"
)

func TestStandardAddress(t *testing.T) {
	for i, test := range []struct {
		input      string
		host, port string
		shouldErr  bool
	}{
		{`localhost`, "localhost", "", false},
		{`localhost:1234`, "localhost", "1234", false},
		{`localhost:`, "localhost", "", false},
		{`0.0.0.0`, "0.0.0.0", "", false},
		{`127.0.0.1:1234`, "127.0.0.1", "1234", false},
		{`:1234`, "", "1234", false},
		{`[::1]`, "::1", "", false},
		{`[::1]:1234`, "::1", "1234", false},
		{`:`, "", "", false},
		{`localhost:http`, "localhost", "http", false},
		{`localhost:https`, "localhost", "https", false},
		{`:http`, "", "http", false},
		{`:https`, "", "https", false},
		{`http://localhost`, "localhost", "http", false},
		{`https://localhost`, "localhost", "https", false},
		{`http://127.0.0.1`, "127.0.0.1", "http", false},
		{`https://127.0.0.1`, "127.0.0.1", "https", false},
		{`http://[::1]`, "::1", "http", false},
		{`http://localhost:1234`, "localhost", "1234", false},
		{`https://127.0.0.1:1234`, "127.0.0.1", "1234", false},
		{`http://[::1]:1234`, "::1", "1234", false},
		{``, "", "", false},
		{`::1`, "::1", "", true},
		{`localhost::`, "localhost::", "", true},
		{`#$%@`, "#$%@", "", true},
	} {
		host, port, err := standardAddress(test.input)

		if err != nil && !test.shouldErr {
			t.Errorf("Test %d: Expected no error, but had error: %v", i, err)
		}
		if err == nil && test.shouldErr {
			t.Errorf("Test %d: Expected error, but had none", i)
		}

		if host != test.host {
			t.Errorf("Test %d: Expected host '%s', got '%s'", i, test.host, host)
		}

		if port != test.port {
			t.Errorf("Test %d: Expected port '%s', got '%s'", i, test.port, port)
		}
	}
}

func TestParseOneAndImport(t *testing.T) {
	setupParseTests()

	testParseOne := func(input string) (serverBlock, error) {
		p := testParser(input)
		p.Next() // parseOne doesn't call Next() to start, so we must
		err := p.parseOne()
		return p.block, err
	}

	for i, test := range []struct {
		input     string
		shouldErr bool
		addresses []address
		tokens    map[string]int // map of directive name to number of tokens expected
	}{
		{`localhost`, false, []address{
			{"localhost", ""},
		}, map[string]int{}},

		{`localhost
		  dir1`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 1,
		}},

		{`localhost:1234
		  dir1 foo bar`, false, []address{
			{"localhost", "1234"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`localhost {
		    dir1
		  }`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 1,
		}},

		{`localhost:1234 {
		    dir1 foo bar
		    dir2
		  }`, false, []address{
			{"localhost", "1234"},
		}, map[string]int{
			"dir1": 3,
			"dir2": 1,
		}},

		{`http://localhost https://localhost
		  dir1 foo bar`, false, []address{
			{"localhost", "http"},
			{"localhost", "https"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`http://localhost https://localhost {
		    dir1 foo bar
		  }`, false, []address{
			{"localhost", "http"},
			{"localhost", "https"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`http://localhost, https://localhost {
		    dir1 foo bar
		  }`, false, []address{
			{"localhost", "http"},
			{"localhost", "https"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`http://localhost, {
		  }`, true, []address{
			{"localhost", "http"},
		}, map[string]int{}},

		{`host1:80, http://host2.com
		  dir1 foo bar
		  dir2 baz`, false, []address{
			{"host1", "80"},
			{"host2.com", "http"},
		}, map[string]int{
			"dir1": 3,
			"dir2": 2,
		}},

		{`http://host1.com,
		  http://host2.com,
		  https://host3.com`, false, []address{
			{"host1.com", "http"},
			{"host2.com", "http"},
			{"host3.com", "https"},
		}, map[string]int{}},

		{`http://host1.com:1234, https://host2.com
		  dir1 foo {
		    bar baz
		  }
		  dir2`, false, []address{
			{"host1.com", "1234"},
			{"host2.com", "https"},
		}, map[string]int{
			"dir1": 6,
			"dir2": 1,
		}},

		{`127.0.0.1
		  dir1 {
		    bar baz
		  }
		  dir2 {
		    foo bar
		  }`, false, []address{
			{"127.0.0.1", ""},
		}, map[string]int{
			"dir1": 5,
			"dir2": 5,
		}},

		{`127.0.0.1
		  unknown_directive`, true, []address{
			{"127.0.0.1", ""},
		}, map[string]int{}},

		{`localhost
		  dir1 {
		    foo`, true, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 3,
		}},

		{`localhost
		  dir1 {
		  }`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 3,
		}},

		{`localhost
		  dir1 {
		  } }`, true, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 3,
		}},

		{`localhost
		  dir1 {
		    nested {
		      foo
		    }
		  }
		  dir2 foo bar`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 7,
			"dir2": 3,
		}},

		{``, false, []address{}, map[string]int{}},

		{`localhost
		  dir1 arg1
		  import import_test1.txt`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 2,
			"dir2": 3,
			"dir3": 1,
		}},

		{`import import_test2.txt`, false, []address{
			{"host1", ""},
		}, map[string]int{
			"dir1": 1,
			"dir2": 2,
		}},

		{`import import_test1.txt import_test2.txt`, true, []address{}, map[string]int{}},

		{`import not_found.txt`, true, []address{}, map[string]int{}},

		{`""`, false, []address{}, map[string]int{}},

		{``, false, []address{}, map[string]int{}},
	} {
		result, err := testParseOne(test.input)

		if test.shouldErr && err == nil {
			t.Errorf("Test %d: Expected an error, but didn't get one", i)
		}
		if !test.shouldErr && err != nil {
			t.Errorf("Test %d: Expected no error, but got: %v", i, err)
		}

		if len(result.Addresses) != len(test.addresses) {
			t.Errorf("Test %d: Expected %d addresses, got %d",
				i, len(test.addresses), len(result.Addresses))
			continue
		}
		for j, addr := range result.Addresses {
			if addr.Host != test.addresses[j].Host {
				t.Errorf("Test %d, address %d: Expected host to be '%s', but was '%s'",
					i, j, test.addresses[j].Host, addr.Host)
			}
			if addr.Port != test.addresses[j].Port {
				t.Errorf("Test %d, address %d: Expected port to be '%s', but was '%s'",
					i, j, test.addresses[j].Port, addr.Port)
			}
		}

		if len(result.Tokens) != len(test.tokens) {
			t.Errorf("Test %d: Expected %d directives, had %d",
				i, len(test.tokens), len(result.Tokens))
			continue
		}
		for directive, tokens := range result.Tokens {
			if len(tokens) != test.tokens[directive] {
				t.Errorf("Test %d, directive '%s': Expected %d tokens, counted %d",
					i, directive, test.tokens[directive], len(tokens))
				continue
			}
		}
	}
}

func TestParseAll(t *testing.T) {
	setupParseTests()

	for i, test := range []struct {
		input     string
		shouldErr bool
		addresses [][]address // addresses per server block, in order
	}{
		{`localhost`, false, [][]address{
			{{"localhost", ""}},
		}},

		{`localhost:1234`, false, [][]address{
			[]address{{"localhost", "1234"}},
		}},

		{`localhost:1234 {
		  }
		  localhost:2015 {
		  }`, false, [][]address{
			[]address{{"localhost", "1234"}},
			[]address{{"localhost", "2015"}},
		}},

		{`localhost:1234, http://host2`, false, [][]address{
			[]address{{"localhost", "1234"}, {"host2", "http"}},
		}},

		{`localhost:1234, http://host2,`, true, [][]address{}},

		{`http://host1.com, http://host2.com {
		  }
		  https://host3.com, https://host4.com {
		  }`, false, [][]address{
			[]address{{"host1.com", "http"}, {"host2.com", "http"}},
			[]address{{"host3.com", "https"}, {"host4.com", "https"}},
		}},
	} {
		p := testParser(test.input)
		blocks, err := p.parseAll()

		if test.shouldErr && err == nil {
			t.Errorf("Test %d: Expected an error, but didn't get one", i)
		}
		if !test.shouldErr && err != nil {
			t.Errorf("Test %d: Expected no error, but got: %v", i, err)
		}

		if len(blocks) != len(test.addresses) {
			t.Errorf("Test %d: Expected %d server blocks, got %d",
				i, len(test.addresses), len(blocks))
			continue
		}
		for j, block := range blocks {
			if len(block.Addresses) != len(test.addresses[j]) {
				t.Errorf("Test %d: Expected %d addresses in block %d, got %d",
					i, len(test.addresses[j]), j, len(block.Addresses))
				continue
			}
			for k, addr := range block.Addresses {
				if addr.Host != test.addresses[j][k].Host {
					t.Errorf("Test %d, block %d, address %d: Expected host to be '%s', but was '%s'",
						i, j, k, test.addresses[j][k].Host, addr.Host)
				}
				if addr.Port != test.addresses[j][k].Port {
					t.Errorf("Test %d, block %d, address %d: Expected port to be '%s', but was '%s'",
						i, j, k, test.addresses[j][k].Port, addr.Port)
				}
			}
		}
	}
}

func setupParseTests() {
	// Set up some bogus directives for testing
	ValidDirectives = map[string]struct{}{
		"dir1": {},
		"dir2": {},
		"dir3": {},
	}
}

func testParser(input string) parser {
	buf := strings.NewReader(input)
	p := parser{Dispenser: NewDispenser("Test", buf), checkDirectives: true}
	return p
}
<commit_msg>Added test for environment replacement.<commit_after>package parse

import (
	"os"
	"strings"
	"testing"
)

func TestStandardAddress(t *testing.T) {
	for i, test := range []struct {
		input      string
		host, port string
		shouldErr  bool
	}{
		{`localhost`, "localhost", "", false},
		{`localhost:1234`, "localhost", "1234", false},
		{`localhost:`, "localhost", "", false},
		{`0.0.0.0`, "0.0.0.0", "", false},
		{`127.0.0.1:1234`, "127.0.0.1", "1234", false},
		{`:1234`, "", "1234", false},
		{`[::1]`, "::1", "", false},
		{`[::1]:1234`, "::1", "1234", false},
		{`:`, "", "", false},
		{`localhost:http`, "localhost", "http", false},
		{`localhost:https`, "localhost", "https", false},
		{`:http`, "", "http", false},
		{`:https`, "", "https", false},
		{`http://localhost`, "localhost", "http", false},
		{`https://localhost`, "localhost", "https", false},
		{`http://127.0.0.1`, "127.0.0.1", "http", false},
		{`https://127.0.0.1`, "127.0.0.1", "https", false},
		{`http://[::1]`, "::1", "http", false},
		{`http://localhost:1234`, "localhost", "1234", false},
		{`https://127.0.0.1:1234`, "127.0.0.1", "1234", false},
		{`http://[::1]:1234`, "::1", "1234", false},
		{``, "", "", false},
		{`::1`, "::1", "", true},
		{`localhost::`, "localhost::", "", true},
		{`#$%@`, "#$%@", "", true},
	} {
		host, port, err := standardAddress(test.input)

		if err != nil && !test.shouldErr {
			t.Errorf("Test %d: Expected no error, but had error: %v", i, err)
		}
		if err == nil && test.shouldErr {
			t.Errorf("Test %d: Expected error, but had none", i)
		}

		if host != test.host {
			t.Errorf("Test %d: Expected host '%s', got '%s'", i, test.host, host)
		}

		if port != test.port {
			t.Errorf("Test %d: Expected port '%s', got '%s'", i, test.port, port)
		}
	}
}

func TestParseOneAndImport(t *testing.T) {
	setupParseTests()

	testParseOne := func(input string) (serverBlock, error) {
		p := testParser(input)
		p.Next() // parseOne doesn't call Next() to start, so we must
		err := p.parseOne()
		return p.block, err
	}

	for i, test := range []struct {
		input     string
		shouldErr bool
		addresses []address
		tokens    map[string]int // map of directive name to number of tokens expected
	}{
		{`localhost`, false, []address{
			{"localhost", ""},
		}, map[string]int{}},

		{`localhost
		  dir1`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 1,
		}},

		{`localhost:1234
		  dir1 foo bar`, false, []address{
			{"localhost", "1234"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`localhost {
		    dir1
		  }`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 1,
		}},

		{`localhost:1234 {
		    dir1 foo bar
		    dir2
		  }`, false, []address{
			{"localhost", "1234"},
		}, map[string]int{
			"dir1": 3,
			"dir2": 1,
		}},

		{`http://localhost https://localhost
		  dir1 foo bar`, false, []address{
			{"localhost", "http"},
			{"localhost", "https"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`http://localhost https://localhost {
		    dir1 foo bar
		  }`, false, []address{
			{"localhost", "http"},
			{"localhost", "https"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`http://localhost, https://localhost {
		    dir1 foo bar
		  }`, false, []address{
			{"localhost", "http"},
			{"localhost", "https"},
		}, map[string]int{
			"dir1": 3,
		}},

		{`http://localhost, {
		  }`, true, []address{
			{"localhost", "http"},
		}, map[string]int{}},

		{`host1:80, http://host2.com
		  dir1 foo bar
		  dir2 baz`, false, []address{
			{"host1", "80"},
			{"host2.com", "http"},
		}, map[string]int{
			"dir1": 3,
			"dir2": 2,
		}},

		{`http://host1.com,
		  http://host2.com,
		  https://host3.com`, false, []address{
			{"host1.com", "http"},
			{"host2.com", "http"},
			{"host3.com", "https"},
		}, map[string]int{}},

		{`http://host1.com:1234, https://host2.com
		  dir1 foo {
		    bar baz
		  }
		  dir2`, false, []address{
			{"host1.com", "1234"},
			{"host2.com", "https"},
		}, map[string]int{
			"dir1": 6,
			"dir2": 1,
		}},

		{`127.0.0.1
		  dir1 {
		    bar baz
		  }
		  dir2 {
		    foo bar
		  }`, false, []address{
			{"127.0.0.1", ""},
		}, map[string]int{
			"dir1": 5,
			"dir2": 5,
		}},

		{`127.0.0.1
		  unknown_directive`, true, []address{
			{"127.0.0.1", ""},
		}, map[string]int{}},

		{`localhost
		  dir1 {
		    foo`, true, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 3,
		}},

		{`localhost
		  dir1 {
		  }`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 3,
		}},

		{`localhost
		  dir1 {
		  } }`, true, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 3,
		}},

		{`localhost
		  dir1 {
		    nested {
		      foo
		    }
		  }
		  dir2 foo bar`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 7,
			"dir2": 3,
		}},

		{``, false, []address{}, map[string]int{}},

		{`localhost
		  dir1 arg1
		  import import_test1.txt`, false, []address{
			{"localhost", ""},
		}, map[string]int{
			"dir1": 2,
			"dir2": 3,
			"dir3": 1,
		}},

		{`import import_test2.txt`, false, []address{
			{"host1", ""},
		}, map[string]int{
			"dir1": 1,
			"dir2": 2,
		}},

		{`import import_test1.txt import_test2.txt`, true, []address{}, map[string]int{}},

		{`import not_found.txt`, true, []address{}, map[string]int{}},

		{`""`, false, []address{}, map[string]int{}},

		{``, false, []address{}, map[string]int{}},
	} {
		result, err := testParseOne(test.input)

		if test.shouldErr && err == nil {
			t.Errorf("Test %d: Expected an error, but didn't get one", i)
		}
		if !test.shouldErr && err != nil {
			t.Errorf("Test %d: Expected no error, but got: %v", i, err)
		}

		if len(result.Addresses) != len(test.addresses) {
			t.Errorf("Test %d: Expected %d addresses, got %d",
				i, len(test.addresses), len(result.Addresses))
			continue
		}
		for j, addr := range result.Addresses {
			if addr.Host != test.addresses[j].Host {
				t.Errorf("Test %d, address %d: Expected host to be '%s', but was '%s'",
					i, j, test.addresses[j].Host, addr.Host)
			}
			if addr.Port != test.addresses[j].Port {
				t.Errorf("Test %d, address %d: Expected port to be '%s', but was '%s'",
					i, j, test.addresses[j].Port, addr.Port)
			}
		}

		if len(result.Tokens) != len(test.tokens) {
			t.Errorf("Test %d: Expected %d directives, had %d",
				i, len(test.tokens), len(result.Tokens))
			continue
		}
		for directive, tokens := range result.Tokens {
			if len(tokens) != test.tokens[directive] {
				t.Errorf("Test %d, directive '%s': Expected %d tokens, counted %d",
					i, directive, test.tokens[directive], len(tokens))
				continue
			}
		}
	}
}

func TestParseAll(t *testing.T) {
	setupParseTests()

	for i, test := range []struct {
		input     string
		shouldErr bool
		addresses [][]address // addresses per server block, in order
	}{
		{`localhost`, false, [][]address{
			{{"localhost", ""}},
		}},

		{`localhost:1234`, false, [][]address{
			[]address{{"localhost", "1234"}},
		}},

		{`localhost:1234 {
		  }
		  localhost:2015 {
		  }`, false, [][]address{
			[]address{{"localhost", "1234"}},
			[]address{{"localhost", "2015"}},
		}},

		{`localhost:1234, http://host2`, false, [][]address{
			[]address{{"localhost", "1234"}, {"host2", "http"}},
		}},

		{`localhost:1234, http://host2,`, true, [][]address{}},

		{`http://host1.com, http://host2.com {
		  }
		  https://host3.com, https://host4.com {
		  }`, false, [][]address{
			[]address{{"host1.com", "http"}, {"host2.com", "http"}},
			[]address{{"host3.com", "https"}, {"host4.com", "https"}},
		}},
	} {
		p := testParser(test.input)
		blocks, err := p.parseAll()

		if test.shouldErr && err == nil {
			t.Errorf("Test %d: Expected an error, but didn't get one", i)
		}
		if !test.shouldErr && err != nil {
			t.Errorf("Test %d: Expected no error, but got: %v", i, err)
		}

		if len(blocks) != len(test.addresses) {
			t.Errorf("Test %d: Expected %d server blocks, got %d",
				i, len(test.addresses), len(blocks))
			continue
		}
		for j, block := range blocks {
			if len(block.Addresses) != len(test.addresses[j]) {
				t.Errorf("Test %d: Expected %d addresses in block %d, got %d",
					i, len(test.addresses[j]), j, len(block.Addresses))
				continue
			}
			for k, addr := range block.Addresses {
				if addr.Host != test.addresses[j][k].Host {
					t.Errorf("Test %d, block %d, address %d: Expected host to be '%s', but was '%s'",
						i, j, k, test.addresses[j][k].Host, addr.Host)
				}
				if addr.Port != test.addresses[j][k].Port {
					t.Errorf("Test %d, block %d, address %d: Expected port to be '%s', but was '%s'",
						i, j, k, test.addresses[j][k].Port, addr.Port)
				}
			}
		}
	}
}

func TestEnvironmentReplacement(t *testing.T) {
	setupParseTests()

	os.Setenv("MY_PORT", "8080")
	os.Setenv("MY_ADDRESS", "servername.com")
	os.Setenv("MY_ADDRESS2", "127.0.0.1")

	for i, test := range []struct {
		input     string
		addresses [][]address // addresses per server block, in order
	}{
		{`{$MY_ADDRESS}`, [][]address{
			{{"servername.com", ""}},
		}},

		{`{$MY_ADDRESS}:{$MY_PORT}`, [][]address{
			[]address{{"servername.com", "8080"}},
		}},

		{`{$MY_ADDRESS2}:1234 {
		  }
		  localhost:{$MY_PORT} {
		  }`, [][]address{
			[]address{{"127.0.0.1", "1234"}},
			[]address{{"localhost", "8080"}},
		}},
	} {
		p := testParser(test.input)
		blocks, err := p.parseAll()

		if err != nil {
			t.Errorf("Test %d: Expected no error, but got: %v", i, err)
		}

		if len(blocks) != len(test.addresses) {
			t.Errorf("Test %d: Expected %d server blocks, got %d",
				i, len(test.addresses), len(blocks))
			continue
		}
		for j, block := range blocks {
			if len(block.Addresses) != len(test.addresses[j]) {
				t.Errorf("Test %d: Expected %d addresses in block %d, got %d",
					i, len(test.addresses[j]), j, len(block.Addresses))
				continue
			}
			for k, addr := range block.Addresses {
				if addr.Host != test.addresses[j][k].Host {
					t.Errorf("Test %d, block %d, address %d: Expected host to be '%s', but was '%s'",
						i, j, k, test.addresses[j][k].Host, addr.Host)
				}
				if addr.Port != test.addresses[j][k].Port {
					t.Errorf("Test %d, block %d, address %d: Expected port to be '%s', but was '%s'",
						i, j, k, test.addresses[j][k].Port, addr.Port)
				}
			}
		}
	}
}

func setupParseTests() {
	// Set up some bogus directives for testing
	ValidDirectives = map[string]struct{}{
		"dir1": {},
		"dir2": {},
		"dir3": {},
	}
}

func testParser(input string) parser {
	buf := strings.NewReader(input)
	p := parser{Dispenser: NewDispenser("Test", buf), checkDirectives: true}
	return p
}
<|endoftext|>