<commit_before>/*
Copyright IBM Corp. 2016 All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

		 http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package system_chaincode

import (
	"fmt"
	"net"
	"os"
	"testing"
	"time"

	"github.com/hyperledger/fabric/core/chaincode"
	"github.com/hyperledger/fabric/core/ledger"
	"github.com/hyperledger/fabric/core/system_chaincode/api"
	"github.com/hyperledger/fabric/core/system_chaincode/samplesyscc"
	"github.com/hyperledger/fabric/core/util"
	pb "github.com/hyperledger/fabric/protos"
	"github.com/spf13/viper"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

// Invoke or query a chaincode.
func invoke(ctx context.Context, spec *pb.ChaincodeSpec, typ pb.Transaction_Type) (*pb.ChaincodeEvent, string, []byte, error) {
	chaincodeInvocationSpec := &pb.ChaincodeInvocationSpec{ChaincodeSpec: spec}

	// Now create the Transactions message and send to Peer.
	uuid := util.GenerateUUID()

	var transaction *pb.Transaction
	var err error
	transaction, err = pb.NewChaincodeExecute(chaincodeInvocationSpec, uuid, typ)
	if err != nil {
		return nil, uuid, nil, fmt.Errorf("Error invoking chaincode: %s ", err)
	}

	var retval []byte
	var execErr error
	var ccevt *pb.ChaincodeEvent
	if typ == pb.Transaction_CHAINCODE_QUERY {
		retval, ccevt, execErr = chaincode.Execute(ctx, chaincode.GetChain(chaincode.DefaultChain), transaction)
	} else {
		ledger, _ := ledger.GetLedger()
		ledger.BeginTxBatch("1")
		retval, ccevt, execErr = chaincode.Execute(ctx, chaincode.GetChain(chaincode.DefaultChain), transaction)
		if err != nil {
			return nil, uuid, nil, fmt.Errorf("Error invoking chaincode: %s ", err)
		}
		ledger.CommitTxBatch("1", []*pb.Transaction{transaction}, nil, nil)
	}

	return ccevt, uuid, retval, execErr
}

func closeListenerAndSleep(l net.Listener) {
	if l != nil {
		l.Close()
		time.Sleep(2 * time.Second)
	}
}

// Test deploy of a transaction.
func TestExecuteDeploySysChaincode(t *testing.T) {
	var opts []grpc.ServerOption
	grpcServer := grpc.NewServer(opts...)
	viper.Set("peer.fileSystemPath", "/var/hyperledger/test/tmpdb")

	//use a different address than what we usually use for "peer"
	//we override the peerAddress set in chaincode_support.go
	peerAddress := "0.0.0.0:40303"
	lis, err := net.Listen("tcp", peerAddress)
	if err != nil {
		t.Fail()
		t.Logf("Error starting peer listener %s", err)
		return
	}

	getPeerEndpoint := func() (*pb.PeerEndpoint, error) {
		return &pb.PeerEndpoint{ID: &pb.PeerID{Name: "testpeer"}, Address: peerAddress}, nil
	}

	ccStartupTimeout := time.Duration(5000) * time.Millisecond
	pb.RegisterChaincodeSupportServer(grpcServer, chaincode.NewChaincodeSupport(chaincode.DefaultChain, getPeerEndpoint, false, ccStartupTimeout, nil))

	go grpcServer.Serve(lis)

	var ctxt = context.Background()

	//set systemChaincodes to sample
	systemChaincodes = []*api.SystemChaincode{
		{
			Enabled:   true,
			Name:      "sample_syscc",
			Path:      "github.com/hyperledger/fabric/core/system_chaincode/samplesyscc",
			InitArgs:  []string{},
			Chaincode: &samplesyscc.SampleSysCC{},
		},
	}

	RegisterSysCCs()

	url := "github.com/hyperledger/fabric/core/system_chaincode/sample_syscc"
	f := "putval"
	args := []string{"greeting", "hey there"}

	spec := &pb.ChaincodeSpec{Type: 1, ChaincodeID: &pb.ChaincodeID{Name: "sample_syscc", Path: url}, CtorMsg: &pb.ChaincodeInput{Function: f, Args: args}}
	_, _, _, err = invoke(ctxt, spec, pb.Transaction_CHAINCODE_INVOKE)
	if err != nil {
		closeListenerAndSleep(lis)
		t.Fail()
		t.Logf("Error invoking sample_syscc: %s", err)
		return
	}

	f = "getval"
	args = []string{"greeting"}
	spec = &pb.ChaincodeSpec{Type: 1, ChaincodeID: &pb.ChaincodeID{Name: "sample_syscc", Path: url}, CtorMsg: &pb.ChaincodeInput{Function: f, Args: args}}
	_, _, _, err = invoke(ctxt, spec, pb.Transaction_CHAINCODE_QUERY)
	if err != nil {
		closeListenerAndSleep(lis)
		t.Fail()
		t.Logf("Error invoking sample_syscc: %s", err)
		return
	}

	cds := &pb.ChaincodeDeploymentSpec{ExecEnv: 1, ChaincodeSpec: &pb.ChaincodeSpec{Type: 1, ChaincodeID: &pb.ChaincodeID{Name: "sample_syscc", Path: url}, CtorMsg: &pb.ChaincodeInput{Args: args}}}

	chaincode.GetChain(chaincode.DefaultChain).Stop(ctxt, cds)

	closeListenerAndSleep(lis)
}

func TestMain(m *testing.M) {
	SetupTestConfig()
	os.Exit(m.Run())
}
<commit_msg>give a unique listening port for system chaincode tests<commit_after>/*
Copyright IBM Corp. 2016 All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

		 http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package system_chaincode

import (
	"fmt"
	"net"
	"os"
	"testing"
	"time"

	"github.com/hyperledger/fabric/core/chaincode"
	"github.com/hyperledger/fabric/core/ledger"
	"github.com/hyperledger/fabric/core/system_chaincode/api"
	"github.com/hyperledger/fabric/core/system_chaincode/samplesyscc"
	"github.com/hyperledger/fabric/core/util"
	pb "github.com/hyperledger/fabric/protos"
	"github.com/spf13/viper"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

// Invoke or query a chaincode.
func invoke(ctx context.Context, spec *pb.ChaincodeSpec, typ pb.Transaction_Type) (*pb.ChaincodeEvent, string, []byte, error) {
	chaincodeInvocationSpec := &pb.ChaincodeInvocationSpec{ChaincodeSpec: spec}

	// Now create the Transactions message and send to Peer.
	uuid := util.GenerateUUID()

	var transaction *pb.Transaction
	var err error
	transaction, err = pb.NewChaincodeExecute(chaincodeInvocationSpec, uuid, typ)
	if err != nil {
		return nil, uuid, nil, fmt.Errorf("Error invoking chaincode: %s ", err)
	}

	var retval []byte
	var execErr error
	var ccevt *pb.ChaincodeEvent
	if typ == pb.Transaction_CHAINCODE_QUERY {
		retval, ccevt, execErr = chaincode.Execute(ctx, chaincode.GetChain(chaincode.DefaultChain), transaction)
	} else {
		ledger, _ := ledger.GetLedger()
		ledger.BeginTxBatch("1")
		retval, ccevt, execErr = chaincode.Execute(ctx, chaincode.GetChain(chaincode.DefaultChain), transaction)
		if err != nil {
			return nil, uuid, nil, fmt.Errorf("Error invoking chaincode: %s ", err)
		}
		ledger.CommitTxBatch("1", []*pb.Transaction{transaction}, nil, nil)
	}

	return ccevt, uuid, retval, execErr
}

func closeListenerAndSleep(l net.Listener) {
	if l != nil {
		l.Close()
		time.Sleep(2 * time.Second)
	}
}

// Test deploy of a transaction.
func TestExecuteDeploySysChaincode(t *testing.T) {
	var opts []grpc.ServerOption
	grpcServer := grpc.NewServer(opts...)
	viper.Set("peer.fileSystemPath", "/var/hyperledger/test/tmpdb")

	//use a different address than what we usually use for "peer"
	//we override the peerAddress set in chaincode_support.go
	peerAddress := "0.0.0.0:41726"
	lis, err := net.Listen("tcp", peerAddress)
	if err != nil {
		t.Fail()
		t.Logf("Error starting peer listener %s", err)
		return
	}

	getPeerEndpoint := func() (*pb.PeerEndpoint, error) {
		return &pb.PeerEndpoint{ID: &pb.PeerID{Name: "testpeer"}, Address: peerAddress}, nil
	}

	ccStartupTimeout := time.Duration(5000) * time.Millisecond
	pb.RegisterChaincodeSupportServer(grpcServer, chaincode.NewChaincodeSupport(chaincode.DefaultChain, getPeerEndpoint, false, ccStartupTimeout, nil))

	go grpcServer.Serve(lis)

	var ctxt = context.Background()

	//set systemChaincodes to sample
	systemChaincodes = []*api.SystemChaincode{
		{
			Enabled:   true,
			Name:      "sample_syscc",
			Path:      "github.com/hyperledger/fabric/core/system_chaincode/samplesyscc",
			InitArgs:  []string{},
			Chaincode: &samplesyscc.SampleSysCC{},
		},
	}

	RegisterSysCCs()

	url := "github.com/hyperledger/fabric/core/system_chaincode/sample_syscc"
	f := "putval"
	args := []string{"greeting", "hey there"}

	spec := &pb.ChaincodeSpec{Type: 1, ChaincodeID: &pb.ChaincodeID{Name: "sample_syscc", Path: url}, CtorMsg: &pb.ChaincodeInput{Function: f, Args: args}}
	_, _, _, err = invoke(ctxt, spec, pb.Transaction_CHAINCODE_INVOKE)
	if err != nil {
		closeListenerAndSleep(lis)
		t.Fail()
		t.Logf("Error invoking sample_syscc: %s", err)
		return
	}

	f = "getval"
	args = []string{"greeting"}
	spec = &pb.ChaincodeSpec{Type: 1, ChaincodeID: &pb.ChaincodeID{Name: "sample_syscc", Path: url}, CtorMsg: &pb.ChaincodeInput{Function: f, Args: args}}
	_, _, _, err = invoke(ctxt, spec, pb.Transaction_CHAINCODE_QUERY)
	if err != nil {
		closeListenerAndSleep(lis)
		t.Fail()
		t.Logf("Error invoking sample_syscc: %s", err)
		return
	}

	cds := &pb.ChaincodeDeploymentSpec{ExecEnv: 1, ChaincodeSpec: &pb.ChaincodeSpec{Type: 1, ChaincodeID: &pb.ChaincodeID{Name: "sample_syscc", Path: url}, CtorMsg: &pb.ChaincodeInput{Args: args}}}

	chaincode.GetChain(chaincode.DefaultChain).Stop(ctxt, cds)

	closeListenerAndSleep(lis)
}

func TestMain(m *testing.M) {
	SetupTestConfig()
	os.Exit(m.Run())
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
)

type ImageFile struct {
	outputWriter io.Writer
	nbytes       int
	bytes        []byte
	wptr         int
	dptr         int
}

var startDIR = flag.Int("dirstart", 16, "Starting sector of directory extent")
var nDIR = flag.Int("ndir", 8, "Number of sectors in the directory extent")
var nSectors = flag.Int("sectors", 2048, "Number of sectors in volume")
var bytesPerSector = flag.Int("bpsect", 512, "Bytes per sector")
var volumeName = flag.String("vol", "WorkDisk", "Volume name of the image")
var outputName = flag.String("out", "sdimage.bin", "Output file containing filesystem image")
var impName = flag.String("import", "imgs", "Directory containing files to import into image")

func main() {
	flag.Parse()

	fout, err := os.Create(*outputName)
	if err != nil {
		log.Fatal(err)
	}
	defer fout.Close()

	f := &ImageFile{}
	f.outputWriter = fout
	f.nbytes = *nSectors * *bytesPerSector
	f.bytes = make([]byte, f.nbytes)
	f.dptr = *startDIR * *bytesPerSector
	for i := 0; i < f.nbytes; i++ {
		f.bytes[i] = 0xCC
	}

	if *nDIR < 1 {
		log.Fatal("Number of directory sectors must at least be 1.")
	}
	endDIR := *startDIR + *nDIR - 1
	f.bind("$DIR", *startDIR, endDIR, 1)
	f.bind(*volumeName, 0, 0, 2)

	startIPL, endIPL := f.place("sys/$IPL")
	startSYS, endSYS := f.place("sys/$SYS")

	fmt.Println("startIPL ", startIPL, "endIPL ", endIPL)
	fmt.Println("startSYS ", startSYS, "endSYS ", endSYS)
	fmt.Println("startDIR ", *startDIR, "endDIR ", endDIR)

	f.wptr = *bytesPerSector * (endDIR+1)
	filesAndDirs, err := ioutil.ReadDir(*impName)
	if err != nil {
		log.Fatal(err)
	}
	for _, fi := range filesAndDirs {
		if fi.IsDir() {
			continue
		}
		fmt.Println(*impName + fi.Name())
		start, end := f.place(*impName + fi.Name())
		fmt.Println("  start ", start, "end ", end)
	}

	r := bytes.NewReader(f.bytes)
	io.Copy(f.outputWriter, r)
}

func (f *ImageFile) place(filename string) (int, int) {
	sectorStart := f.wptr / 512

	file, err := os.Open(filename)
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()

	size, err := file.Seek(0, os.SEEK_END)
	if err != nil {
		log.Fatal(err)
	}
	file.Seek(0, os.SEEK_SET)
	size = (size + 511) & -512

	bs, err := ioutil.ReadAll(file)
	if err != nil {
		log.Fatal(err)
	}

	copy(f.bytes[f.wptr:], bs)
	f.wptr = f.wptr + int(size)

	sectorEnd := (f.wptr - 1) / 512

	basename := path.Base(filename)
	f.bind(basename, sectorStart, sectorEnd, 1)
	return sectorStart, sectorEnd
}

func (f *ImageFile) bind(filename string, start, end, kind int) {
	if len(filename) > 47 {
		filename = filename[0:46]
	}
	f.bytes[f.dptr] = byte(len(filename))
	copy(f.bytes[f.dptr+1:], filename)
	f.bytes[f.dptr+48] = byte(kind)
	f.bytes[f.dptr+49] = 0
	f.bytes[f.dptr+50] = byte(start & 255)
	f.bytes[f.dptr+51] = byte((start >> 8) & 255)
	f.bytes[f.dptr+52] = byte(end & 255)
	f.bytes[f.dptr+53] = byte((end >> 8) & 255)
	for i := 54; i < 64; i++ {
		f.bytes[f.dptr+i] = 0
	}
	f.dptr = f.dptr + 64
}
<commit_msg>Fix annoying bug wrt pathnames<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
)

type ImageFile struct {
	outputWriter io.Writer
	nbytes       int
	bytes        []byte
	wptr         int
	dptr         int
}

var startDIR = flag.Int("dirstart", 16, "Starting sector of directory extent")
var nDIR = flag.Int("ndir", 8, "Number of sectors in the directory extent")
var nSectors = flag.Int("sectors", 2048, "Number of sectors in volume")
var bytesPerSector = flag.Int("bpsect", 512, "Bytes per sector")
var volumeName = flag.String("vol", "WorkDisk", "Volume name of the image")
var outputName = flag.String("out", "sdimage.bin", "Output file containing filesystem image")
var impName = flag.String("import", "imgs", "Directory containing files to import into image.")

func main() {
	flag.Parse()

	fout, err := os.Create(*outputName)
	if err != nil {
		log.Fatal(err)
	}
	defer fout.Close()

	f := &ImageFile{}
	f.outputWriter = fout
	f.nbytes = *nSectors * *bytesPerSector
	f.bytes = make([]byte, f.nbytes)
	f.dptr = *startDIR * *bytesPerSector
	for i := 0; i < f.nbytes; i++ {
		f.bytes[i] = 0xCC
	}

	if *nDIR < 1 {
		log.Fatal("Number of directory sectors must at least be 1.")
	}
	endDIR := *startDIR + *nDIR - 1
	f.bind("$DIR", *startDIR, endDIR, 1)
	f.bind(*volumeName, 0, 0, 2)

	startIPL, endIPL := f.place("sys/$IPL")
	startSYS, endSYS := f.place("sys/$SYS")

	fmt.Println("startIPL ", startIPL, "endIPL ", endIPL)
	fmt.Println("startSYS ", startSYS, "endSYS ", endSYS)
	fmt.Println("startDIR ", *startDIR, "endDIR ", endDIR)

	f.wptr = *bytesPerSector * (endDIR+1)
	filesAndDirs, err := ioutil.ReadDir(*impName)
	if err != nil {
		log.Fatal(err)
	}
	for _, fi := range filesAndDirs {
		if fi.IsDir() {
			continue
		}
		fmt.Println(*impName + "/" + fi.Name())
		start, end := f.place(*impName + "/" + fi.Name())
		fmt.Println("  start ", start, "end ", end)
	}

	r := bytes.NewReader(f.bytes)
	io.Copy(f.outputWriter, r)
}

func (f *ImageFile) place(filename string) (int, int) {
	sectorStart := f.wptr / 512

	file, err := os.Open(filename)
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()

	size, err := file.Seek(0, os.SEEK_END)
	if err != nil {
		log.Fatal(err)
	}
	file.Seek(0, os.SEEK_SET)
	size = (size + 511) & -512

	bs, err := ioutil.ReadAll(file)
	if err != nil {
		log.Fatal(err)
	}

	copy(f.bytes[f.wptr:], bs)
	f.wptr = f.wptr + int(size)

	sectorEnd := (f.wptr - 1) / 512

	basename := path.Base(filename)
	f.bind(basename, sectorStart, sectorEnd, 1)
	return sectorStart, sectorEnd
}

func (f *ImageFile) bind(filename string, start, end, kind int) {
	if len(filename) > 47 {
		filename = filename[0:46]
	}
	f.bytes[f.dptr] = byte(len(filename))
	copy(f.bytes[f.dptr+1:], filename)
	f.bytes[f.dptr+48] = byte(kind)
	f.bytes[f.dptr+49] = 0
	f.bytes[f.dptr+50] = byte(start & 255)
	f.bytes[f.dptr+51] = byte((start >> 8) & 255)
	f.bytes[f.dptr+52] = byte(end & 255)
	f.bytes[f.dptr+53] = byte((end >> 8) & 255)
	for i := 54; i < 64; i++ {
		f.bytes[f.dptr+i] = 0
	}
	f.dptr = f.dptr + 64
}
<|endoftext|>


<commit_before>package multitenant

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"flag"
	"math"
	"strings"
	"sync"
	"time"

	log "github.com/sirupsen/logrus"
	billing "github.com/weaveworks/billing-client"

	"github.com/weaveworks/scope/app"
	"github.com/weaveworks/scope/report"
)

// BillingEmitterConfig has everything we need to make a billing emitter
type BillingEmitterConfig struct {
	Enabled         bool
	DefaultInterval time.Duration
	UserIDer        UserIDer
}

// RegisterFlags registers the billing emitter flags with the main flag set.
func (cfg *BillingEmitterConfig) RegisterFlags(f *flag.FlagSet) {
	f.BoolVar(&cfg.Enabled, "app.billing.enabled", false, "enable emitting billing info")
	f.DurationVar(&cfg.DefaultInterval, "app.billing.default-publish-interval", 3*time.Second, "default publish interval to assume for reports")
}

// BillingEmitter is the billing emitter
type BillingEmitter struct {
	app.Collector
	BillingEmitterConfig
	billing *billing.Client

	sync.Mutex
	intervalCache map[string]time.Duration
	rounding      map[string]float64
}

// NewBillingEmitter changes a new billing emitter which emits billing events
func NewBillingEmitter(upstream app.Collector, billingClient *billing.Client, cfg BillingEmitterConfig) (*BillingEmitter, error) {
	return &BillingEmitter{
		Collector:            upstream,
		billing:              billingClient,
		BillingEmitterConfig: cfg,
		intervalCache:        make(map[string]time.Duration),
		rounding:             make(map[string]float64),
	}, nil
}

// Add implements app.Collector
func (e *BillingEmitter) Add(ctx context.Context, rep report.Report, buf []byte) error {
	now := time.Now().UTC()
	userID, err := e.UserIDer(ctx)
	if err != nil {
		// Underlying collector needs to get userID too, so it's OK to abort
		// here. If this fails, so will underlying collector so no point
		// proceeding.
		return err
	}
	rowKey, colKey := calculateDynamoKeys(userID, now)

	interval := e.reportInterval(rep)
	// Cache the last-known value of interval for this user, and use
	// it if we didn't find one in this report.
	e.Lock()
	if interval != 0 {
		e.intervalCache[userID] = interval
	} else {
		if lastKnown, found := e.intervalCache[userID]; found {
			interval = lastKnown
		} else {
			interval = e.DefaultInterval
		}
	}
	// Billing takes an integer number of seconds, so keep track of the amount lost to rounding
	nodeSeconds := interval.Seconds()*float64(len(rep.Host.Nodes)) + e.rounding[userID]
	rounding := nodeSeconds - math.Floor(nodeSeconds)
	e.rounding[userID] = rounding
	e.Unlock()

	hasher := sha256.New()
	hasher.Write(buf)
	hash := "sha256:" + base64.URLEncoding.EncodeToString(hasher.Sum(nil))

	weaveNetCount := 0
	if hasWeaveNet(rep) {
		weaveNetCount = 1
	}

	amounts := billing.Amounts{
		billing.ContainerSeconds: int64(interval/time.Second) * int64(len(rep.Container.Nodes)),
		billing.NodeSeconds:      int64(nodeSeconds),
		billing.WeaveNetSeconds:  int64(interval/time.Second) * int64(weaveNetCount),
	}
	metadata := map[string]string{
		"row_key": rowKey,
		"col_key": colKey,
	}

	err = e.billing.AddAmounts(
		hash,
		userID,
		now,
		amounts,
		metadata,
	)
	if err != nil {
		// No return, because we want to proceed even if we fail to emit
		// billing data, so that defects in the billing system don't break
		// report collection. Just log the fact & carry on.
		log.Errorf("Failed emitting billing data: %v", err)
	}

	return e.Collector.Add(ctx, rep, buf)
}

func commandParameter(cmd, flag string) (string, bool) {
	i := strings.Index(cmd, flag)
	if i != -1 {
		// here we expect the command looks like `-foo=bar` or `-foo bar`
		aft := strings.Fields(cmd[i+len(flag):])
		if len(aft) > 0 && len(aft[0]) > 0 {
			if aft[0][0] == '=' {
				return aft[0][1:], true
			}
			return aft[0], true
		}
	}
	return "", false
}

func intervalFromCommand(cmd string) string {
	if strings.Contains(cmd, "scope") {
		if publishInterval, ok := commandParameter(cmd, "probe.publish.interval"); ok {
			// If spy interval is higher than publish interval, some reports will have no process data
			if spyInterval, ok := commandParameter(cmd, "spy.interval"); ok {
				pubDuration, err1 := time.ParseDuration(publishInterval)
				spyDuration, err2 := time.ParseDuration(spyInterval)
				if err1 == nil && err2 == nil && spyDuration > pubDuration {
					return spyInterval
				}
			}
			return publishInterval
		}
	}
	return ""
}

// reportInterval tries to find the custom report interval of this report. If
// it is malformed, or not set, it returns zero.
func (e *BillingEmitter) reportInterval(r report.Report) time.Duration {
	if r.Window != 0 {
		return r.Window
	}
	var inter string
	for _, c := range r.Container.Nodes {
		if cmd, ok := c.Latest.Lookup(report.DockerContainerCommand); ok {
			if inter = intervalFromCommand(cmd); inter != "" {
				break
			}
		}
	}
	if inter == "" { // not found in containers: look in processes
		for _, c := range r.Process.Nodes {
			if cmd, ok := c.Latest.Lookup(report.Cmdline); ok {
				if inter = intervalFromCommand(cmd); inter != "" {
					break
				}
			}
		}
	}
	if inter == "" {
		return 0
	}
	d, err := time.ParseDuration(inter)
	if err != nil {
		return 0
	}
	return d
}

// Tries to determine if this report came from a host running Weave Net
func hasWeaveNet(r report.Report) bool {
	for _, n := range r.Overlay.Nodes {
		overlayType, _ := report.ParseOverlayNodeID(n.ID)
		if overlayType == report.WeaveOverlayPeerPrefix {
			return true
		}
	}

	return false
}

// Close shuts down the billing emitter and billing client flushing events.
func (e *BillingEmitter) Close() {
	e.Collector.Close()
	_ = e.billing.Close()
}
<commit_msg>multitenant: only count real hosts for billing<commit_after>package multitenant

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"flag"
	"math"
	"strings"
	"sync"
	"time"

	log "github.com/sirupsen/logrus"
	billing "github.com/weaveworks/billing-client"

	"github.com/weaveworks/scope/app"
	"github.com/weaveworks/scope/report"
)

// BillingEmitterConfig has everything we need to make a billing emitter
type BillingEmitterConfig struct {
	Enabled         bool
	DefaultInterval time.Duration
	UserIDer        UserIDer
}

// RegisterFlags registers the billing emitter flags with the main flag set.
func (cfg *BillingEmitterConfig) RegisterFlags(f *flag.FlagSet) {
	f.BoolVar(&cfg.Enabled, "app.billing.enabled", false, "enable emitting billing info")
	f.DurationVar(&cfg.DefaultInterval, "app.billing.default-publish-interval", 3*time.Second, "default publish interval to assume for reports")
}

// BillingEmitter is the billing emitter
type BillingEmitter struct {
	app.Collector
	BillingEmitterConfig
	billing *billing.Client

	sync.Mutex
	intervalCache map[string]time.Duration
	rounding      map[string]float64
}

// NewBillingEmitter changes a new billing emitter which emits billing events
func NewBillingEmitter(upstream app.Collector, billingClient *billing.Client, cfg BillingEmitterConfig) (*BillingEmitter, error) {
	return &BillingEmitter{
		Collector:            upstream,
		billing:              billingClient,
		BillingEmitterConfig: cfg,
		intervalCache:        make(map[string]time.Duration),
		rounding:             make(map[string]float64),
	}, nil
}

// Add implements app.Collector
func (e *BillingEmitter) Add(ctx context.Context, rep report.Report, buf []byte) error {
	now := time.Now().UTC()
	userID, err := e.UserIDer(ctx)
	if err != nil {
		// Underlying collector needs to get userID too, so it's OK to abort
		// here. If this fails, so will underlying collector so no point
		// proceeding.
		return err
	}
	rowKey, colKey := calculateDynamoKeys(userID, now)

	interval, nodes := e.scanReport(rep)
	// Cache the last-known value of interval for this user, and use
	// it if we didn't find one in this report.
	e.Lock()
	if interval != 0 {
		e.intervalCache[userID] = interval
	} else {
		if lastKnown, found := e.intervalCache[userID]; found {
			interval = lastKnown
		} else {
			interval = e.DefaultInterval
		}
	}
	// Billing takes an integer number of seconds, so keep track of the amount lost to rounding
	nodeSeconds := interval.Seconds()*float64(nodes) + e.rounding[userID]
	rounding := nodeSeconds - math.Floor(nodeSeconds)
	e.rounding[userID] = rounding
	e.Unlock()

	hasher := sha256.New()
	hasher.Write(buf)
	hash := "sha256:" + base64.URLEncoding.EncodeToString(hasher.Sum(nil))

	weaveNetCount := 0
	if hasWeaveNet(rep) {
		weaveNetCount = 1
	}

	amounts := billing.Amounts{
		billing.ContainerSeconds: int64(interval/time.Second) * int64(len(rep.Container.Nodes)),
		billing.NodeSeconds:      int64(nodeSeconds),
		billing.WeaveNetSeconds:  int64(interval/time.Second) * int64(weaveNetCount),
	}
	metadata := map[string]string{
		"row_key": rowKey,
		"col_key": colKey,
	}

	err = e.billing.AddAmounts(
		hash,
		userID,
		now,
		amounts,
		metadata,
	)
	if err != nil {
		// No return, because we want to proceed even if we fail to emit
		// billing data, so that defects in the billing system don't break
		// report collection. Just log the fact & carry on.
		log.Errorf("Failed emitting billing data: %v", err)
	}

	return e.Collector.Add(ctx, rep, buf)
}

func commandParameter(cmd, flag string) (string, bool) {
	i := strings.Index(cmd, flag)
	if i != -1 {
		// here we expect the command looks like `-foo=bar` or `-foo bar`
		aft := strings.Fields(cmd[i+len(flag):])
		if len(aft) > 0 && len(aft[0]) > 0 {
			if aft[0][0] == '=' {
				return aft[0][1:], true
			}
			return aft[0], true
		}
	}
	return "", false
}

func intervalFromCommand(cmd string) string {
	if strings.Contains(cmd, "scope") {
		if publishInterval, ok := commandParameter(cmd, "probe.publish.interval"); ok {
			// If spy interval is higher than publish interval, some reports will have no process data
			if spyInterval, ok := commandParameter(cmd, "spy.interval"); ok {
				pubDuration, err1 := time.ParseDuration(publishInterval)
				spyDuration, err2 := time.ParseDuration(spyInterval)
				if err1 == nil && err2 == nil && spyDuration > pubDuration {
					return spyInterval
				}
			}
			return publishInterval
		}
	}
	return ""
}

// scanReport counts the nodes tries to find any custom report interval
// of this report. If it is malformed, or not set, it returns zero.
func (e *BillingEmitter) scanReport(r report.Report) (time.Duration, int) {
	nHosts := 0
	// We scan the host nodes looking for ones reported by a per-node probe;
	// the Kubernetes cluster probe also makes host nodes but they only have a few fields set
	for _, h := range r.Host.Nodes {
		// Relying here on Uptime being something that changes in each report, hence will be in a delta report
		if _, ok := h.Latest.Lookup(report.Uptime); ok {
			nHosts++
		}
	}
	if r.Window != 0 {
		return r.Window, nHosts
	}
	var inter string
	for _, c := range r.Container.Nodes {
		if cmd, ok := c.Latest.Lookup(report.DockerContainerCommand); ok {
			if inter = intervalFromCommand(cmd); inter != "" {
				break
			}
		}
	}
	if inter == "" { // not found in containers: look in processes
		for _, c := range r.Process.Nodes {
			if cmd, ok := c.Latest.Lookup(report.Cmdline); ok {
				if inter = intervalFromCommand(cmd); inter != "" {
					break
				}
			}
		}
	}
	if inter == "" {
		return 0, nHosts
	}
	d, err := time.ParseDuration(inter)
	if err != nil {
		return 0, nHosts
	}
	return d, nHosts
}

// Tries to determine if this report came from a host running Weave Net
func hasWeaveNet(r report.Report) bool {
	for _, n := range r.Overlay.Nodes {
		overlayType, _ := report.ParseOverlayNodeID(n.ID)
		if overlayType == report.WeaveOverlayPeerPrefix {
			return true
		}
	}

	return false
}

// Close shuts down the billing emitter and billing client flushing events.
func (e *BillingEmitter) Close() {
	e.Collector.Close()
	_ = e.billing.Close()
}
<|endoftext|>


<commit_before>package udptransport

import (
	"io"
	"net"
	"sync"

	"github.com/juju/errors"
	"github.com/peer-calls/peer-calls/server/logger"
	"github.com/peer-calls/peer-calls/server/servertransport"
	"github.com/peer-calls/peer-calls/server/stringmux"
	"github.com/peer-calls/peer-calls/server/udpmux"
)

// Manager is in charge of managing server-to-server UDP Transports. The
// overarching design is as follows.
//
//  1. UDPMux is used to demultiplex UDP packets coming in from different Peer
//     Calls nodes based on remote addr.
//  2. For each incoming server packet from a specific remote address, a new
//     transport factory is created. A transport factory can also be created
//     manually.
//  3. Each factory creates a separate transport peer room, and it uses the
//     stringmux package to figure out which packets are for which room.
//  4. Each stream transport then uses a stringmux again to figure out which
//     packet is for which transport component:
//     - packets with 'm' prefix are media packets for MediaTransport, and
//     - packets with 's' prefix are for SCTP component which is used for
//       DataTransport and MetadataTransport.
//
// Due to the issues with sctp connection closure, it might be wise to create
// long-lived SCTP connection per factory and demultiplex packets separately.
type Manager struct {
	params *ManagerParams

	// udpMux is used for demultiplexing UDP packets from other server nodes.
	udpMux *udpmux.UDPMux

	// torndown will be closed when manager is closed.
	torndown chan struct{}

	// factoriesChan contains accepted Factories.
	factoriesChan chan *Factory
	closeOnce     sync.Once
	mu            sync.RWMutex
	wg            sync.WaitGroup

	// factories is the map of all created and active Factories.
	factories map[*stringmux.StringMux]*Factory
}

// ManagerParams are the parameters for Manager.
type ManagerParams struct {
	// Conn is the packet connection to use for sending server-to-server data.
	Conn net.PacketConn
	Log  logger.Logger
}

// NewManager creates a new instance of Manager.
func NewManager(params ManagerParams) *Manager {
	params.Log = params.Log.WithNamespaceAppended("transport_manager")

	readChanSize := 100

	udpMux := udpmux.New(udpmux.Params{
		Conn:           params.Conn,
		MTU:            uint32(servertransport.ReceiveMTU),
		Log:            params.Log,
		ReadChanSize:   readChanSize,
		ReadBufferSize: 0,
	})

	t := &Manager{
		params:        &params,
		udpMux:        udpMux,
		torndown:      make(chan struct{}),
		factoriesChan: make(chan *Factory),
		factories:     make(map[*stringmux.StringMux]*Factory),
	}

	t.wg.Add(1)

	go func() {
		defer t.wg.Done()
		t.start()
	}()

	return t
}

func (t *Manager) Factories() []*Factory {
	t.mu.RLock()
	defer t.mu.RUnlock()

	factories := make([]*Factory, 0, len(t.factories))

	for _, factory := range t.factories {
		factories = append(factories, factory)
	}

	return factories
}

func (t *Manager) start() {
	for {
		conn, err := t.udpMux.AcceptConn()
		if err != nil {
			t.params.Log.Error("Accept UDPMux conn", errors.Trace(err), nil)

			return
		}

		log := t.params.Log.WithCtx(logger.Ctx{
			"remote_addr": conn.RemoteAddr(),
		})

		log.Info("Accept UDP conn", nil)

		factory, err := t.createFactory(conn)
		if err != nil {
			t.params.Log.Error("Create Transport Factory", errors.Trace(err), nil)

			return
		}

		t.factoriesChan <- factory
	}
}

// createFactory creates a new Factory for the provided
// connection.
func (t *Manager) createFactory(conn net.Conn) (*Factory, error) {
	t.mu.Lock()
	defer t.mu.Unlock()

	readChanSize := 100

	stringMux := stringmux.New(stringmux.Params{
		Log:            t.params.Log,
		Conn:           conn,
		MTU:            uint32(servertransport.ReceiveMTU), // TODO not sure if this is ok
		ReadChanSize:   readChanSize,
		ReadBufferSize: 0,
	})

	factory := NewFactory(t.params.Log, &t.wg, stringMux)
	t.factories[stringMux] = factory

	t.wg.Add(1)

	go func() {
		defer t.wg.Done()
		<-stringMux.Done()

		t.mu.Lock()
		defer t.mu.Unlock()

		delete(t.factories, stringMux)
	}()

	return factory, nil
}

func (t *Manager) AcceptFactory() (*Factory, error) {
	factory, ok := <-t.factoriesChan
	if !ok {
		return nil, errors.Annotate(io.ErrClosedPipe, "Manager is tearing down")
	}

	return factory, nil
}

func (t *Manager) GetFactory(raddr net.Addr) (*Factory, error) {
	conn, err := t.udpMux.GetConn(raddr)
	if err != nil {
		return nil, errors.Annotatef(err, "getting conn for raddr: %s", raddr)
	}

	return t.createFactory(conn)
}

func (t *Manager) Close() error {
	err := t.close()

	t.wg.Wait()

	return err
}

func (t *Manager) Done() <-chan struct{} {
	return t.torndown
}

func (t *Manager) close() error {
	t.mu.Lock()
	defer t.mu.Unlock()

	err := t.udpMux.Close()

	t.closeOnce.Do(func() {
		close(t.factoriesChan)

		for stringMux, factory := range t.factories {
			_ = stringMux.Close()

			factory.Close()

			delete(t.factories, stringMux)
		}

		close(t.torndown)
	})

	return errors.Trace(err)
}
<commit_msg>Add a TODO comment to udptransport.Manager<commit_after>package udptransport

import (
	"io"
	"net"
	"sync"

	"github.com/juju/errors"
	"github.com/peer-calls/peer-calls/server/logger"
	"github.com/peer-calls/peer-calls/server/servertransport"
	"github.com/peer-calls/peer-calls/server/stringmux"
	"github.com/peer-calls/peer-calls/server/udpmux"
)

// Manager is in charge of managing server-to-server UDP Transports. The
// overarching design is as follows.
//
//  1. UDPMux is used to demultiplex UDP packets coming in from different Peer
//     Calls nodes based on remote addr.
//  2. For each incoming server packet from a specific remote address, a new
//     transport factory is created. A transport factory can also be created
//     manually.
//  3. Each factory creates a separate transport peer room, and it uses the
//     stringmux package to figure out which packets are for which room.
//  4. Each stream transport then uses a stringmux again to figure out which
//     packet is for which transport component:
//     - packets with 'm' prefix are media packets for MediaTransport, and
//     - packets with 's' prefix are for SCTP component which is used for
//       DataTransport and MetadataTransport.
//
// TODO Due to the issues with sctp connection closure, it might be wise to
// create long-lived SCTP connection per factory and demultiplex packets
// separately. To clarify, the steps modified so that:
//
// 1. stringmux conns with 'm' and 's' prefixes are created in (2). sctp
//    association for 's' is created when the factory is created.
// 2. The stringmux package will be used twice to determine:
//    - Which SCTP stream packet should go to which (room) DataTransport and
//      MetadataTransport.
//    - Which Media packet should go to which MediaTransport
//
// The above should allow for the use of a single, long-lived SCTP association
// between two Peer Calls nodes.
//
// NOTE: I'm not sure about the performance issues this might have, but it's
// the apparent solution to issues with caused by terminating SCTP associations
// without a abort or shutdown signals.
type Manager struct {
	params *ManagerParams

	// udpMux is used for demultiplexing UDP packets from other server nodes.
	udpMux *udpmux.UDPMux

	// torndown will be closed when manager is closed.
	torndown chan struct{}

	// factoriesChan contains accepted Factories.
	factoriesChan chan *Factory
	closeOnce     sync.Once
	mu            sync.RWMutex
	wg            sync.WaitGroup

	// factories is the map of all created and active Factories.
	factories map[*stringmux.StringMux]*Factory
}

// ManagerParams are the parameters for Manager.
type ManagerParams struct {
	// Conn is the packet connection to use for sending server-to-server data.
	Conn net.PacketConn
	Log  logger.Logger
}

// NewManager creates a new instance of Manager.
func NewManager(params ManagerParams) *Manager {
	params.Log = params.Log.WithNamespaceAppended("transport_manager")

	readChanSize := 100

	udpMux := udpmux.New(udpmux.Params{
		Conn:           params.Conn,
		MTU:            uint32(servertransport.ReceiveMTU),
		Log:            params.Log,
		ReadChanSize:   readChanSize,
		ReadBufferSize: 0,
	})

	t := &Manager{
		params:        &params,
		udpMux:        udpMux,
		torndown:      make(chan struct{}),
		factoriesChan: make(chan *Factory),
		factories:     make(map[*stringmux.StringMux]*Factory),
	}

	t.wg.Add(1)

	go func() {
		defer t.wg.Done()
		t.start()
	}()

	return t
}

func (t *Manager) Factories() []*Factory {
	t.mu.RLock()
	defer t.mu.RUnlock()

	factories := make([]*Factory, 0, len(t.factories))

	for _, factory := range t.factories {
		factories = append(factories, factory)
	}

	return factories
}

func (t *Manager) start() {
	for {
		conn, err := t.udpMux.AcceptConn()
		if err != nil {
			t.params.Log.Error("Accept UDPMux conn", errors.Trace(err), nil)

			return
		}

		log := t.params.Log.WithCtx(logger.Ctx{
			"remote_addr": conn.RemoteAddr(),
		})

		log.Info("Accept UDP conn", nil)

		factory, err := t.createFactory(conn)
		if err != nil {
			t.params.Log.Error("Create Transport Factory", errors.Trace(err), nil)

			return
		}

		t.factoriesChan <- factory
	}
}

// createFactory creates a new Factory for the provided
// connection.
func (t *Manager) createFactory(conn net.Conn) (*Factory, error) {
	t.mu.Lock()
	defer t.mu.Unlock()

	readChanSize := 100

	stringMux := stringmux.New(stringmux.Params{
		Log:            t.params.Log,
		Conn:           conn,
		MTU:            uint32(servertransport.ReceiveMTU), // TODO not sure if this is ok
		ReadChanSize:   readChanSize,
		ReadBufferSize: 0,
	})

	factory := NewFactory(t.params.Log, &t.wg, stringMux)
	t.factories[stringMux] = factory

	t.wg.Add(1)

	go func() {
		defer t.wg.Done()
		<-stringMux.Done()

		t.mu.Lock()
		defer t.mu.Unlock()

		delete(t.factories, stringMux)
	}()

	return factory, nil
}

func (t *Manager) AcceptFactory() (*Factory, error) {
	factory, ok := <-t.factoriesChan
	if !ok {
		return nil, errors.Annotate(io.ErrClosedPipe, "Manager is tearing down")
	}

	return factory, nil
}

func (t *Manager) GetFactory(raddr net.Addr) (*Factory, error) {
	conn, err := t.udpMux.GetConn(raddr)
	if err != nil {
		return nil, errors.Annotatef(err, "getting conn for raddr: %s", raddr)
	}

	return t.createFactory(conn)
}

func (t *Manager) Close() error {
	err := t.close()

	t.wg.Wait()

	return err
}

func (t *Manager) Done() <-chan struct{} {
	return t.torndown
}

func (t *Manager) close() error {
	t.mu.Lock()
	defer t.mu.Unlock()

	err := t.udpMux.Close()

	t.closeOnce.Do(func() {
		close(t.factoriesChan)

		for stringMux, factory := range t.factories {
			_ = stringMux.Close()

			factory.Close()

			delete(t.factories, stringMux)
		}

		close(t.torndown)
	})

	return errors.Trace(err)
}
<|endoftext|>


<commit_before>package service

import (
	"encoding/json"
	"github.com/zwirec/TGChatScanner/clarifaiApi"
	"github.com/zwirec/TGChatScanner/modelManager"
	"github.com/zwirec/TGChatScanner/requestHandler"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"os/user"
	"sync"
	"syscall"
	"github.com/zwirec/TGChatScanner/TGBotApi"
)

type Config map[string]map[string]interface{}

//Service s
type Service struct {
	mux         *http.ServeMux
	srv         *http.Server
	rAPIHandler *requestHandler.RequestHandler
	config      Config
	logger      *log.Logger
}

func NewService() *Service {
	return &Service{
		rAPIHandler: requestHandler.NewRequestHandler(),
		mux:         http.NewServeMux(),
		logger:      log.New(os.Stdout, "", log.LstdFlags),
	}
}

func (s *Service) Run() error {
	configUrl := os.Getenv("TGCHATSCANNER_REMOTE_CONFIG")
	rc := true

	if configUrl == "" {
		s.logger.Println("Using local config")

		rc = false
		usr, err := user.Current()

		if err != nil {
			s.logger.Println(err)
			return err
		}

		configUrl = usr.HomeDir + "/.config/tgchatscanner/config.json"
	} else {
		s.logger.Println("Using remote config")
	}

	if err := s.parseConfig(configUrl, rc); err != nil {
		s.logger.Println(err)
		return err
	}

	s.signalProcessing()

	db, err := modelManager.ConnectToDB(s.config["db"])
	if err != nil {
		s.logger.Println(err)
		return err
	}

	clApi := clarifaiApi.NewClarifaiApi(s.config["clarifai"]["api_key"].(string))

	botApi := TGBotApi.NewBotApi(s.config["tg_bot_api"]["token"].(string))

	workers_n, ok := s.config["server"]["workers"].(int)

	if !ok {
		workers_n = 10
	}
	fdp := requestHandler.NewFileDownloaderPool(workers_n, 100)

	php := requestHandler.NewPhotoHandlersPool(10, 100)

	cache := requestHandler.MemoryCache{}
	context := requestHandler.AppContext{
		Db:            db,
		Downloaders:   fdp,
		PhotoHandlers: php,
		BotApi:        botApi,
		CfApi:         clApi,
		Cache:         &cache,
		Logger:        s.logger,
	}

	s.rAPIHandler.SetAppContext(&context)
	s.rAPIHandler.RegisterHandlers()

	s.srv = &http.Server{Handler: s.rAPIHandler}

	var wg sync.WaitGroup

	wg.Add(1)

	go func() {
		defer wg.Done()

		l, err := net.Listen("unix", s.config["server"]["socket"].(string))
		if err != nil {
			s.logger.Println(err)
			wg.Done()
		}
		s.logger.Println("Socket opened")
		defer os.Remove(s.config["server"]["socket"].(string))
		defer l.Close()

		log.Println("Server started")
		if err := s.srv.Serve(l); err != nil {
			s.logger.Println(err)
			wg.Done()
		}
	}()

	wg.Wait()
	return nil
}

func (s *Service) parseConfig(url string, remote bool) error {
	var configRaw []byte

	if remote {
		res, err := http.Get(url)
		if err != nil {
			s.logger.Println(err)
			return err
		}

		configRaw, err = ioutil.ReadAll(res.Body)
		res.Body.Close()
		if err != nil {
			s.logger.Println(err)
			return err
		}

	} else {
		var err error

		configRaw, err = ioutil.ReadFile(url)
		if err != nil {
			s.logger.Println(err)
			return err
		}
	}

	if err := json.Unmarshal(configRaw, &s.config); err != nil {
		s.logger.Println(err)
		return err
	}

	return nil

}

func (s *Service) signalProcessing() {
	c := make(chan os.Signal)
	signal.Notify(c, syscall.SIGINT)
	go s.handler(c)
}

func (s *Service) handler(c chan os.Signal) {
	for {
		<-c
		log.Print("Gracefully stopping...")
		s.srv.Shutdown(nil)
		os.Exit(0)
	}
}
<commit_msg>Add socket file permissions setting<commit_after>package service

import (
	"encoding/json"
	"github.com/zwirec/TGChatScanner/TGBotApi"
	"github.com/zwirec/TGChatScanner/clarifaiApi"
	"github.com/zwirec/TGChatScanner/modelManager"
	"github.com/zwirec/TGChatScanner/requestHandler"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"os/user"
	"sync"
	"syscall"
)

type Config map[string]map[string]interface{}

//Service s
type Service struct {
	mux         *http.ServeMux
	srv         *http.Server
	rAPIHandler *requestHandler.RequestHandler
	config      Config
	logger      *log.Logger
}

func NewService() *Service {
	return &Service{
		rAPIHandler: requestHandler.NewRequestHandler(),
		mux:         http.NewServeMux(),
		logger:      log.New(os.Stdout, "", log.LstdFlags),
	}
}

func (s *Service) Run() error {
	configUrl := os.Getenv("TGCHATSCANNER_REMOTE_CONFIG")
	rc := true

	if configUrl == "" {
		s.logger.Println("Using local config")

		rc = false
		usr, err := user.Current()

		if err != nil {
			s.logger.Println(err)
			return err
		}

		configUrl = usr.HomeDir + "/.config/tgchatscanner/config.json"
	} else {
		s.logger.Println("Using remote config")
	}

	if err := s.parseConfig(configUrl, rc); err != nil {
		s.logger.Println(err)
		return err
	}

	s.signalProcessing()

	db, err := modelManager.ConnectToDB(s.config["db"])
	if err != nil {
		s.logger.Println(err)
		return err
	}

	clApi := clarifaiApi.NewClarifaiApi(s.config["clarifai"]["api_key"].(string))

	botApi := TGBotApi.NewBotApi(s.config["tg_bot_api"]["token"].(string))

	workers_n, ok := s.config["server"]["workers"].(int)

	if !ok {
		workers_n = 10
	}
	fdp := requestHandler.NewFileDownloaderPool(workers_n, 100)

	php := requestHandler.NewPhotoHandlersPool(10, 100)

	cache := requestHandler.MemoryCache{}
	context := requestHandler.AppContext{
		Db:            db,
		Downloaders:   fdp,
		PhotoHandlers: php,
		BotApi:        botApi,
		CfApi:         clApi,
		Cache:         &cache,
		Logger:        s.logger,
	}

	s.rAPIHandler.SetAppContext(&context)
	s.rAPIHandler.RegisterHandlers()

	s.srv = &http.Server{Handler: s.rAPIHandler}

	var wg sync.WaitGroup

	wg.Add(1)

	go func() {
		defer wg.Done()

		l, err := net.Listen("unix", s.config["server"]["socket"].(string))
		if err != nil {
			s.logger.Println(err)
			wg.Done()
		}

		if err := os.Chmod(s.config["server"]["socket"].(string), 0777); err != nil {
			s.logger.Println(err)
			wg.Done()
		}

		s.logger.Println("Socket opened")
		defer os.Remove(s.config["server"]["socket"].(string))
		defer l.Close()

		log.Println("Server started")
		if err := s.srv.Serve(l); err != nil {
			s.logger.Println(err)
			wg.Done()
		}
	}()

	wg.Wait()
	return nil
}

func (s *Service) parseConfig(url string, remote bool) error {
	var configRaw []byte

	if remote {
		res, err := http.Get(url)
		if err != nil {
			s.logger.Println(err)
			return err
		}

		configRaw, err = ioutil.ReadAll(res.Body)
		res.Body.Close()
		if err != nil {
			s.logger.Println(err)
			return err
		}

	} else {
		var err error

		configRaw, err = ioutil.ReadFile(url)
		if err != nil {
			s.logger.Println(err)
			return err
		}
	}

	if err := json.Unmarshal(configRaw, &s.config); err != nil {
		s.logger.Println(err)
		return err
	}

	return nil

}

func (s *Service) signalProcessing() {
	c := make(chan os.Signal)
	signal.Notify(c, syscall.SIGINT)
	go s.handler(c)
}

func (s *Service) handler(c chan os.Signal) {
	for {
		<-c
		log.Print("Gracefully stopping...")
		s.srv.Shutdown(nil)
		os.Exit(0)
	}
}
<|endoftext|>


<commit_before>package session

import (
	"encoding/gob"
	"errors"
	"net/http"

	"github.com/RichardKnop/recall/config"
	"github.com/gorilla/sessions"
)

// Service wraps session functionality
type Service struct {
	sessionStore   sessions.Store
	sessionOptions *sessions.Options
	session        *sessions.Session
	r              *http.Request
	w              http.ResponseWriter
}

// UserSession has user data stored in a session after logging in
type UserSession struct {
	ClientID     string
	Username     string
	AccessToken  string
	RefreshToken string
}

const (
	storageSessionName  = "recall_session"
	userSessionKey      = "recall_user"
)

var (
	errSessonNotStarted = errors.New("Session not started")
)

func init() {
	// Register a new datatype for storage in sessions
	gob.Register(new(UserSession))
}

// NewService starts a new Service instance
func NewService(cnf *config.Config, r *http.Request, w http.ResponseWriter) *Service {
	return &Service{
		// Session cookie storage
		sessionStore: sessions.NewCookieStore([]byte(cnf.Session.Secret)),
		// Session options
		sessionOptions: &sessions.Options{
			Path:     cnf.Session.Path,
			MaxAge:   cnf.Session.MaxAge,
			HttpOnly: cnf.Session.HTTPOnly,
		},
		r: r,
		w: w,
	}
}

// StartSession starts a new session. This method must be called before other
// public methods of this struct as it sets the internal session object
func (s *Service) StartSession() error {
	session, err := s.sessionStore.Get(s.r, storageSessionName)
	if err != nil {
		return err
	}
	s.session = session
	return nil
}

// GetUserSession returns the user session
func (s *Service) GetUserSession() (*UserSession, error) {
	// Make sure StartSession has been called
	if s.session == nil {
		return nil, errSessonNotStarted
	}

	// Retrieve our user session struct and type-assert it
	userSession, ok := s.session.Values[userSessionKey].(*UserSession)
	if !ok {
		return nil, errors.New("User session type assertion error")
	}

	return userSession, nil
}

// SetUserSession saves the user session
func (s *Service) SetUserSession(userSession *UserSession) error {
	// Make sure StartSession has been called
	if s.session == nil {
		return errSessonNotStarted
	}

	// Set a new user session
	s.session.Values[userSessionKey] = userSession
	return s.session.Save(s.r, s.w)
}

// ClearUserSession deletes the user session
func (s *Service) ClearUserSession() error {
	// Make sure StartSession has been called
	if s.session == nil {
		return errSessonNotStarted
	}

	// Delete the user session
	delete(s.session.Values, userSessionKey)
	return s.session.Save(s.r, s.w)
}

// SetFlashMessage sets a flash message,
// useful for displaying an error after 302 redirection
func (s *Service) SetFlashMessage(msg string) error {
	// Make sure StartSession has been called
	if s.session == nil {
		return errSessonNotStarted
	}

	// Add the flash message
	s.session.AddFlash(msg)
	return s.session.Save(s.r, s.w)
}

// GetFlashMessage returns the first flash message
func (s *Service) GetFlashMessage() (interface{}, error) {
	// Make sure StartSession has been called
	if s.session == nil {
		return nil, errSessonNotStarted
	}

	// Get the last flash message from the stack
	if flashes := s.session.Flashes(); len(flashes) > 0 {
		// We need to save the session, otherwise the flash message won't be removed
		s.session.Save(s.r, s.w)
		return flashes[0], nil
	}

	// No flash messages in the stack
	return nil, nil
}
<commit_msg>Fixed a session bug introduced in the previous commit.<commit_after>package session

import (
	"encoding/gob"
	"errors"
	"net/http"

	"github.com/RichardKnop/recall/config"
	"github.com/gorilla/sessions"
)

// Service wraps session functionality
type Service struct {
	sessionStore   sessions.Store
	sessionOptions *sessions.Options
	session        *sessions.Session
	r              *http.Request
	w              http.ResponseWriter
}

// UserSession has user data stored in a session after logging in
type UserSession struct {
	ClientID     string
	Username     string
	AccessToken  string
	RefreshToken string
}

var (
	storageSessionName  = "recall_session"
	userSessionKey      = "recall_user"
	errSessonNotStarted = errors.New("Session not started")
)

func init() {
	// Register a new datatype for storage in sessions
	gob.Register(new(UserSession))
}

// NewService starts a new Service instance
func NewService(cnf *config.Config, r *http.Request, w http.ResponseWriter) *Service {
	return &Service{
		// Session cookie storage
		sessionStore: sessions.NewCookieStore([]byte(cnf.Session.Secret)),
		// Session options
		sessionOptions: &sessions.Options{
			Path:     cnf.Session.Path,
			MaxAge:   cnf.Session.MaxAge,
			HttpOnly: cnf.Session.HTTPOnly,
		},
		r: r,
		w: w,
	}
}

// StartSession starts a new session. This method must be called before other
// public methods of this struct as it sets the internal session object
func (s *Service) StartSession() error {
	session, err := s.sessionStore.Get(s.r, storageSessionName)
	if err != nil {
		return err
	}
	s.session = session
	return nil
}

// GetUserSession returns the user session
func (s *Service) GetUserSession() (*UserSession, error) {
	// Make sure StartSession has been called
	if s.session == nil {
		return nil, errSessonNotStarted
	}

	// Retrieve our user session struct and type-assert it
	userSession, ok := s.session.Values[userSessionKey].(*UserSession)
	if !ok {
		return nil, errors.New("User session type assertion error")
	}

	return userSession, nil
}

// SetUserSession saves the user session
func (s *Service) SetUserSession(userSession *UserSession) error {
	// Make sure StartSession has been called
	if s.session == nil {
		return errSessonNotStarted
	}

	// Set a new user session
	s.session.Values[userSessionKey] = userSession
	return s.session.Save(s.r, s.w)
}

// ClearUserSession deletes the user session
func (s *Service) ClearUserSession() error {
	// Make sure StartSession has been called
	if s.session == nil {
		return errSessonNotStarted
	}

	// Delete the user session
	delete(s.session.Values, userSessionKey)
	return s.session.Save(s.r, s.w)
}

// SetFlashMessage sets a flash message,
// useful for displaying an error after 302 redirection
func (s *Service) SetFlashMessage(msg string) error {
	// Make sure StartSession has been called
	if s.session == nil {
		return errSessonNotStarted
	}

	// Add the flash message
	s.session.AddFlash(msg)
	return s.session.Save(s.r, s.w)
}

// GetFlashMessage returns the first flash message
func (s *Service) GetFlashMessage() (interface{}, error) {
	// Make sure StartSession has been called
	if s.session == nil {
		return nil, errSessonNotStarted
	}

	// Get the last flash message from the stack
	if flashes := s.session.Flashes(); len(flashes) > 0 {
		// We need to save the session, otherwise the flash message won't be removed
		s.session.Save(s.r, s.w)
		return flashes[0], nil
	}

	// No flash messages in the stack
	return nil, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"reflect"
	"strings"
	"testing"
	"time"
)

func Test_parseCGIHeaders(t *testing.T) {
	data := []struct {
		in      string
		out     string
		headers map[string]string
	}{
		{
			in:      "Some text",
			out:     "Some text",
			headers: map[string]string{},
		},
		{
			in:      "Location: url\n\nSome text",
			out:     "Some text",
			headers: map[string]string{"Location": "url"},
		},
		{
			in:      "Location: url\n\n",
			out:     "",
			headers: map[string]string{"Location": "url"},
		},
		{
			in:      "Location: url\nX-Name:  x-value\n\nSome text",
			out:     "Some text",
			headers: map[string]string{"Location": "url", "X-Name": "x-value"},
		},
		{
			in:      "Some text\nText\n\ntext",
			out:     "Some text\nText\n\ntext",
			headers: map[string]string{},
		},
		{
			in:      "Some text\nText: value in text\n\ntext",
			out:     "Some text\nText: value in text\n\ntext",
			headers: map[string]string{},
		},
		{
			in:      "Text::::\n\ntext",
			out:     "text",
			headers: map[string]string{"Text": ":::"},
		},
		{
			in:      "Text:     :::\n\ntext",
			out:     "text",
			headers: map[string]string{"Text": ":::"},
		},
		{
			in:      "Text:     \n\ntext",
			out:     "Text:     \n\ntext",
			headers: map[string]string{},
		},
		{
			in:      "Header: value\nText:     \n\ntext",
			out:     "Header: value\nText:     \n\ntext",
			headers: map[string]string{},
		},
	}

	for i, item := range data {
		out, headers := parseCGIHeaders(item.in)
		if !reflect.DeepEqual(item.headers, headers) || item.out != out {
			t.Errorf("%d:\nexpected: %s / %#v\nreal    : %s / %#v", i, item.out, item.headers, out, headers)
		}
	}
}

func Test_getShellAndParams(t *testing.T) {
	shell, params, err := getShellAndParams("ls", Config{shell: "sh", defaultShell: "sh", defaultShOpt: "-c"})
	if shell != "sh" || !reflect.DeepEqual(params, []string{"-c", "ls"}) || err != nil {
		t.Errorf("1. getShellAndParams() failed")
	}

	shell, params, err = getShellAndParams("ls", Config{shell: "bash", defaultShell: "sh", defaultShOpt: "-c"})
	if shell != "bash" || !reflect.DeepEqual(params, []string{"-c", "ls"}) || err != nil {
		t.Errorf("3. getShellAndParams() failed")
	}

	shell, params, err = getShellAndParams("ls -l -a", Config{shell: "", defaultShell: "sh", defaultShOpt: "-c"})
	if shell != "ls" || !reflect.DeepEqual(params, []string{"-l", "-a"}) || err != nil {
		t.Errorf("4. getShellAndParams() failed")
	}

	shell, params, err = getShellAndParams("ls -l 'a b'", Config{shell: "", defaultShell: "sh", defaultShOpt: "-c"})
	if shell != "ls" || !reflect.DeepEqual(params, []string{"-l", "a b"}) || err != nil {
		t.Errorf("5. getShellAndParams() failed")
	}

	_, _, err = getShellAndParams("ls '-l", Config{shell: "", defaultShell: "sh", defaultShOpt: "-c"})
	if err == nil {
		t.Errorf("6. getShellAndParams() failed")
	}
}

func Test_getShellAndParams_windows(t *testing.T) {
	shell, params, err := getShellAndParams("ls", Config{shell: "cmd", defaultShell: "cmd", defaultShOpt: "/C"})
	if shell != "cmd" || !reflect.DeepEqual(params, []string{"/C", "ls"}) || err != nil {
		t.Errorf("2. getShellAndParams() failed")
	}
}

func httpRequest(method string, url string, postData string) ([]byte, error) {
	var postDataReader io.Reader
	if method == "POST" && len(postData) > 0 {
		postDataReader = strings.NewReader(postData)
	}

	request, err := http.NewRequest(method, url, postDataReader)
	if err != nil {
		return nil, err
	}

	request.Header.Set("X-Real-Ip", "127.0.0.1")
	client := &http.Client{}
	res, err := client.Do(request)
	if err != nil {
		return nil, err
	}

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

	err = res.Body.Close()
	if err != nil {
		return nil, err
	}

	return body, nil
}

func getFreePort(t *testing.T) string {
	listen, _ := net.Listen("tcp", ":0")
	parts := strings.Split(listen.Addr().String(), ":")
	err := listen.Close()
	if err != nil {
		t.Errorf("getFreePort() failed")
	}

	return parts[len(parts)-1]
}

func testHTTP(t *testing.T, method, url, postData string, fn func(body string) bool, message string) {
	res, err := httpRequest(method, url, postData)
	if err != nil {
		t.Errorf("%s, get %s failed: %s", message, url, err)
	}
	if !fn(string(res)) {
		t.Errorf("%s failed", message)
	}
}

func Test_main(t *testing.T) {
	port := getFreePort(t)
	os.Args = []string{"shell2http",
		"-add-exit",
		"-cache=1",
		"-cgi",
		// "-export-all-vars",
		"-export-vars=HOME",
		"-one-thread",
		"-shell=",
		"-log=/dev/null",
		"-port=" + port,
		"GET:/echo", "echo 123",
		"POST:/form", "echo var=$v_var",
		"/error", "/ not exists cmd",
		"POST:/post", "cat",
		"/redirect", `echo "Location: /` + "\n" + `"`,
	}
	go main()
	time.Sleep(100 * time.Millisecond) // wait for up http server

	// hide stderr
	oldStderr := os.Stderr // keep backup of the real stderr
	newStderr, err := os.Open("/dev/null")
	if err != nil {
		t.Errorf("open /dev/null: %s", err)
	}
	os.Stderr = newStderr
	defer func() {
		os.Stderr = oldStderr
		err := newStderr.Close()
		if err != nil {
			t.Errorf("Stderr Close failed: %s", err)
		}
	}()

	testHTTP(t, "GET", "http://localhost:"+port+"/", "",
		func(res string) bool { return len(res) > 0 && strings.HasPrefix(res, "<!DOCTYPE html>") },
		"1. get /",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/echo", "",
		func(res string) bool { return res == "123\n" },
		"2. echo",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/echo", "",
		func(res string) bool { return res == "123\n" },
		"3. echo from cache",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/404", "",
		func(res string) bool { return strings.HasPrefix(res, "404 page not found") },
		"4. 404",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/error", "",
		func(res string) bool { return strings.HasPrefix(res, "exec error:") },
		"5. error",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/redirect", "",
		func(res string) bool { return strings.HasPrefix(res, "<!DOCTYPE html>") },
		"6. redirect",
	)

	testHTTP(t, "POST", "http://localhost:"+port+"/post", "X-header: value\n\ntext",
		func(res string) bool { return strings.HasPrefix(res, "text") },
		"7. POST",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/form", "",
		func(res string) bool {
			return strings.HasPrefix(res, http.StatusText(http.StatusMethodNotAllowed))
		},
		"8. POST with GET",
	)
}

func Test_errChain(t *testing.T) {
	err := errChain()
	if err != nil {
		t.Errorf("1. errChain() empty failed")
	}

	err = errChain(func() error { return nil })
	if err != nil {
		t.Errorf("2. errChain() failed")
	}

	err = errChain(func() error { return nil }, func() error { return nil })
	if err != nil {
		t.Errorf("3. errChain() failed")
	}

	err = errChain(func() error { return fmt.Errorf("error") })
	if err == nil {
		t.Errorf("4. errChain() failed")
	}

	err = errChain(func() error { return nil }, func() error { return fmt.Errorf("error") })
	if err == nil {
		t.Errorf("5. errChain() failed")
	}

	var1 := false
	err = errChain(func() error { return fmt.Errorf("error") }, func() error { var1 = true; return nil })
	if err == nil || var1 {
		t.Errorf("6. errChain() failed")
	}
}

func Test_errChainAll(t *testing.T) {
	err := errChainAll()
	if err != nil {
		t.Errorf("1. errChainAll() empty failed")
	}

	err = errChainAll(func() error { return nil })
	if err != nil {
		t.Errorf("2. errChainAll() failed")
	}

	err = errChainAll(func() error { return nil }, func() error { return nil })
	if err != nil {
		t.Errorf("3. errChainAll() failed")
	}

	err = errChainAll(func() error { return fmt.Errorf("error") })
	if err == nil {
		t.Errorf("4. errChainAll() failed")
	}

	err = errChainAll(func() error { return nil }, func() error { return fmt.Errorf("error") })
	if err == nil {
		t.Errorf("5. errChainAll() failed")
	}

	var1 := false
	err = errChainAll(func() error { return fmt.Errorf("error") }, func() error { var1 = true; return nil })
	if err == nil || !var1 {
		t.Errorf("6. errChainAll() failed")
	}
}

func Test_parsePathAndCommands(t *testing.T) {
	tests := []struct {
		name    string
		args    []string
		want    []Command
		wantErr bool
	}{
		{
			name:    "empty list",
			args:    nil,
			want:    nil,
			wantErr: true,
		},
		{
			name:    "empty list 2",
			args:    []string{},
			want:    nil,
			wantErr: true,
		},
		{
			name:    "one arg",
			args:    []string{"arg"},
			want:    nil,
			wantErr: true,
		},
		{
			name:    "two arg without path",
			args:    []string{"arg", "arg2"},
			want:    nil,
			wantErr: true,
		},
		{
			name:    "three arg",
			args:    []string{"/arg", "date", "aaa"},
			want:    nil,
			wantErr: true,
		},
		{
			name:    "two arg",
			args:    []string{"/date", "date"},
			want:    []Command{{path: "/date", cmd: "date"}},
			wantErr: false,
		},
		{
			name:    "four arg",
			args:    []string{"/date", "date", "/", "echo index"},
			want:    []Command{{path: "/date", cmd: "date"}, {path: "/", cmd: "echo index"}},
			wantErr: false,
		},
		{
			name:    "with http method",
			args:    []string{"POST:/date", "date", "GET:/", "echo index"},
			want:    []Command{{path: "/date", cmd: "date", httpMethod: "POST"}, {path: "/", cmd: "echo index", httpMethod: "GET"}},
			wantErr: false,
		},
		{
			name:    "invalid method",
			args:    []string{"get:/date"},
			want:    nil,
			wantErr: true,
		},
		{
			name:    "invalid method2",
			args:    []string{"GET_A:/date"},
			want:    nil,
			wantErr: true,
		},
		{
			name:    "invalid path",
			args:    []string{"GET:/date 2"},
			want:    nil,
			wantErr: true,
		},
		{
			name:    "not uniq path",
			args:    []string{"POST:/date", "date", "POST:/date", "echo index"},
			want:    nil,
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := parsePathAndCommands(tt.args)
			if (err != nil) != tt.wantErr {
				t.Errorf("parsePathAndCommands() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("parsePathAndCommands() = %v, want %v", got, tt.want)
			}
		})
	}
}
<commit_msg>Refactored test function signature<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"reflect"
	"strings"
	"testing"
	"time"
)

func Test_parseCGIHeaders(t *testing.T) {
	data := []struct {
		in      string
		out     string
		headers map[string]string
	}{
		{
			in:      "Some text",
			out:     "Some text",
			headers: map[string]string{},
		},
		{
			in:      "Location: url\n\nSome text",
			out:     "Some text",
			headers: map[string]string{"Location": "url"},
		},
		{
			in:      "Location: url\n\n",
			out:     "",
			headers: map[string]string{"Location": "url"},
		},
		{
			in:      "Location: url\nX-Name:  x-value\n\nSome text",
			out:     "Some text",
			headers: map[string]string{"Location": "url", "X-Name": "x-value"},
		},
		{
			in:      "Some text\nText\n\ntext",
			out:     "Some text\nText\n\ntext",
			headers: map[string]string{},
		},
		{
			in:      "Some text\nText: value in text\n\ntext",
			out:     "Some text\nText: value in text\n\ntext",
			headers: map[string]string{},
		},
		{
			in:      "Text::::\n\ntext",
			out:     "text",
			headers: map[string]string{"Text": ":::"},
		},
		{
			in:      "Text:     :::\n\ntext",
			out:     "text",
			headers: map[string]string{"Text": ":::"},
		},
		{
			in:      "Text:     \n\ntext",
			out:     "Text:     \n\ntext",
			headers: map[string]string{},
		},
		{
			in:      "Header: value\nText:     \n\ntext",
			out:     "Header: value\nText:     \n\ntext",
			headers: map[string]string{},
		},
	}

	for i, item := range data {
		out, headers := parseCGIHeaders(item.in)
		if !reflect.DeepEqual(item.headers, headers) || item.out != out {
			t.Errorf("%d:\nexpected: %s / %#v\nreal    : %s / %#v", i, item.out, item.headers, out, headers)
		}
	}
}

func Test_getShellAndParams(t *testing.T) {
	shell, params, err := getShellAndParams("ls", Config{shell: "sh", defaultShell: "sh", defaultShOpt: "-c"})
	if shell != "sh" || !reflect.DeepEqual(params, []string{"-c", "ls"}) || err != nil {
		t.Errorf("1. getShellAndParams() failed")
	}

	shell, params, err = getShellAndParams("ls", Config{shell: "bash", defaultShell: "sh", defaultShOpt: "-c"})
	if shell != "bash" || !reflect.DeepEqual(params, []string{"-c", "ls"}) || err != nil {
		t.Errorf("3. getShellAndParams() failed")
	}

	shell, params, err = getShellAndParams("ls -l -a", Config{shell: "", defaultShell: "sh", defaultShOpt: "-c"})
	if shell != "ls" || !reflect.DeepEqual(params, []string{"-l", "-a"}) || err != nil {
		t.Errorf("4. getShellAndParams() failed")
	}

	shell, params, err = getShellAndParams("ls -l 'a b'", Config{shell: "", defaultShell: "sh", defaultShOpt: "-c"})
	if shell != "ls" || !reflect.DeepEqual(params, []string{"-l", "a b"}) || err != nil {
		t.Errorf("5. getShellAndParams() failed")
	}

	_, _, err = getShellAndParams("ls '-l", Config{shell: "", defaultShell: "sh", defaultShOpt: "-c"})
	if err == nil {
		t.Errorf("6. getShellAndParams() failed")
	}
}

func Test_getShellAndParams_windows(t *testing.T) {
	shell, params, err := getShellAndParams("ls", Config{shell: "cmd", defaultShell: "cmd", defaultShOpt: "/C"})
	if shell != "cmd" || !reflect.DeepEqual(params, []string{"/C", "ls"}) || err != nil {
		t.Errorf("2. getShellAndParams() failed")
	}
}

func httpRequest(method, url, postData string) ([]byte, error) {
	var postDataReader io.Reader
	if method == "POST" && len(postData) > 0 {
		postDataReader = strings.NewReader(postData)
	}

	request, err := http.NewRequest(method, url, postDataReader)
	if err != nil {
		return nil, err
	}

	request.Header.Set("X-Real-Ip", "127.0.0.1")
	client := &http.Client{}
	res, err := client.Do(request)
	if err != nil {
		return nil, err
	}

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

	err = res.Body.Close()
	if err != nil {
		return nil, err
	}

	return body, nil
}

func getFreePort(t *testing.T) string {
	listen, _ := net.Listen("tcp", ":0")
	parts := strings.Split(listen.Addr().String(), ":")
	err := listen.Close()
	if err != nil {
		t.Errorf("getFreePort() failed")
	}

	return parts[len(parts)-1]
}

func testHTTP(t *testing.T, method, url, postData string, fn func(body string) bool, message string) {
	res, err := httpRequest(method, url, postData)
	if err != nil {
		t.Errorf("%s, get %s failed: %s", message, url, err)
	}
	if !fn(string(res)) {
		t.Errorf("%s failed", message)
	}
}

func Test_main(t *testing.T) {
	port := getFreePort(t)
	os.Args = []string{"shell2http",
		"-add-exit",
		"-cache=1",
		"-cgi",
		// "-export-all-vars",
		"-export-vars=HOME",
		"-one-thread",
		"-shell=",
		"-log=/dev/null",
		"-port=" + port,
		"GET:/echo", "echo 123",
		"POST:/form", "echo var=$v_var",
		"/error", "/ not exists cmd",
		"POST:/post", "cat",
		"/redirect", `echo "Location: /` + "\n" + `"`,
	}
	go main()
	time.Sleep(100 * time.Millisecond) // wait for up http server

	// hide stderr
	oldStderr := os.Stderr // keep backup of the real stderr
	newStderr, err := os.Open("/dev/null")
	if err != nil {
		t.Errorf("open /dev/null: %s", err)
	}
	os.Stderr = newStderr
	defer func() {
		os.Stderr = oldStderr
		err := newStderr.Close()
		if err != nil {
			t.Errorf("Stderr Close failed: %s", err)
		}
	}()

	testHTTP(t, "GET", "http://localhost:"+port+"/", "",
		func(res string) bool { return len(res) > 0 && strings.HasPrefix(res, "<!DOCTYPE html>") },
		"1. get /",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/echo", "",
		func(res string) bool { return res == "123\n" },
		"2. echo",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/echo", "",
		func(res string) bool { return res == "123\n" },
		"3. echo from cache",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/404", "",
		func(res string) bool { return strings.HasPrefix(res, "404 page not found") },
		"4. 404",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/error", "",
		func(res string) bool { return strings.HasPrefix(res, "exec error:") },
		"5. error",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/redirect", "",
		func(res string) bool { return strings.HasPrefix(res, "<!DOCTYPE html>") },
		"6. redirect",
	)

	testHTTP(t, "POST", "http://localhost:"+port+"/post", "X-header: value\n\ntext",
		func(res string) bool { return strings.HasPrefix(res, "text") },
		"7. POST",
	)

	testHTTP(t, "GET", "http://localhost:"+port+"/form", "",
		func(res string) bool {
			return strings.HasPrefix(res, http.StatusText(http.StatusMethodNotAllowed))
		},
		"8. POST with GET",
	)
}

func Test_errChain(t *testing.T) {
	err := errChain()
	if err != nil {
		t.Errorf("1. errChain() empty failed")
	}

	err = errChain(func() error { return nil })
	if err != nil {
		t.Errorf("2. errChain() failed")
	}

	err = errChain(func() error { return nil }, func() error { return nil })
	if err != nil {
		t.Errorf("3. errChain() failed")
	}

	err = errChain(func() error { return fmt.Errorf("error") })
	if err == nil {
		t.Errorf("4. errChain() failed")
	}

	err = errChain(func() error { return nil }, func() error { return fmt.Errorf("error") })
	if err == nil {
		t.Errorf("5. errChain() failed")
	}

	var1 := false
	err = errChain(func() error { return fmt.Errorf("error") }, func() error { var1 = true; return nil })
	if err == nil || var1 {
		t.Errorf("6. errChain() failed")
	}
}

func Test_errChainAll(t *testing.T) {
	err := errChainAll()
	if err != nil {
		t.Errorf("1. errChainAll() empty failed")
	}

	err = errChainAll(func() error { return nil })
	if err != nil {
		t.Errorf("2. errChainAll() failed")
	}

	err = errChainAll(func() error { return nil }, func() error { return nil })
	if err != nil {
		t.Errorf("3. errChainAll() failed")
	}

	err = errChainAll(func() error { return fmt.Errorf("error") })
	if err == nil {
		t.Errorf("4. errChainAll() failed")
	}

	err = errChainAll(func() error { return nil }, func() error { return fmt.Errorf("error") })
	if err == nil {
		t.Errorf("5. errChainAll() failed")
	}

	var1 := false
	err = errChainAll(func() error { return fmt.Errorf("error") }, func() error { var1 = true; return nil })
	if err == nil || !var1 {
		t.Errorf("6. errChainAll() failed")
	}
}

func Test_parsePathAndCommands(t *testing.T) {
	tests := []struct {
		name    string
		args    []string
		want    []Command
		wantErr bool
	}{
		{
			name:    "empty list",
			args:    nil,
			want:    nil,
			wantErr: true,
		},
		{
			name:    "empty list 2",
			args:    []string{},
			want:    nil,
			wantErr: true,
		},
		{
			name:    "one arg",
			args:    []string{"arg"},
			want:    nil,
			wantErr: true,
		},
		{
			name:    "two arg without path",
			args:    []string{"arg", "arg2"},
			want:    nil,
			wantErr: true,
		},
		{
			name:    "three arg",
			args:    []string{"/arg", "date", "aaa"},
			want:    nil,
			wantErr: true,
		},
		{
			name:    "two arg",
			args:    []string{"/date", "date"},
			want:    []Command{{path: "/date", cmd: "date"}},
			wantErr: false,
		},
		{
			name:    "four arg",
			args:    []string{"/date", "date", "/", "echo index"},
			want:    []Command{{path: "/date", cmd: "date"}, {path: "/", cmd: "echo index"}},
			wantErr: false,
		},
		{
			name:    "with http method",
			args:    []string{"POST:/date", "date", "GET:/", "echo index"},
			want:    []Command{{path: "/date", cmd: "date", httpMethod: "POST"}, {path: "/", cmd: "echo index", httpMethod: "GET"}},
			wantErr: false,
		},
		{
			name:    "invalid method",
			args:    []string{"get:/date"},
			want:    nil,
			wantErr: true,
		},
		{
			name:    "invalid method2",
			args:    []string{"GET_A:/date"},
			want:    nil,
			wantErr: true,
		},
		{
			name:    "invalid path",
			args:    []string{"GET:/date 2"},
			want:    nil,
			wantErr: true,
		},
		{
			name:    "not uniq path",
			args:    []string{"POST:/date", "date", "POST:/date", "echo index"},
			want:    nil,
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := parsePathAndCommands(tt.args)
			if (err != nil) != tt.wantErr {
				t.Errorf("parsePathAndCommands() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("parsePathAndCommands() = %v, want %v", got, tt.want)
			}
		})
	}
}
<|endoftext|>


<commit_before>package xpi // import "go.mozilla.org/autograph/signer/xpi"

import (
	"bytes"
	"crypto"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"fmt"
	"strings"

	"github.com/pkg/errors"
	"go.mozilla.org/autograph/signer"
	"go.mozilla.org/cose"
)

const (
	// algHeaderValue compresses to 1 for the key "alg"
	algHeaderValue = 1
	// kidHeaderValue compresses to 4 for the key "kid"
	kidHeaderValue = 4
)

// stringToCOSEAlg returns the cose.Algorithm for a string or nil if
// the algorithm isn't implemented
func stringToCOSEAlg(s string) (v *cose.Algorithm) {
	switch strings.ToUpper(s) {
	case cose.PS256.Name:
		v = cose.PS256
	case cose.ES256.Name:
		v = cose.ES256
	case cose.ES384.Name:
		v = cose.ES384
	case cose.ES512.Name:
		v = cose.ES512
	default:
		v = nil
	}
	return v
}

// stringToCOSEAlg returns the cose.Algorithm for an int or nil if
// the algorithm isn't implemented
func intToCOSEAlg(i int) (v *cose.Algorithm) {
	switch i {
	case cose.PS256.Value:
		v = cose.PS256
	case cose.ES256.Value:
		v = cose.ES256
	case cose.ES384.Value:
		v = cose.ES384
	case cose.ES512.Value:
		v = cose.ES512
	default:
		v = nil
	}
	return v
}

// generateIssuerEEKeyPair returns a public and private key pair for
// the provided COSEAlgorithm
func (s *XPISigner) generateCOSEKeyPair(coseAlg *cose.Algorithm) (eeKey crypto.PrivateKey, eePublicKey crypto.PublicKey, err error) {
	var signer *cose.Signer

	switch coseAlg {
	case nil:
		err = errors.New("Cannot generate private key for nil cose Algorithm")
	case cose.PS256:
		if s.issuerKey == nil {
			err = errors.New("Cannot generate COSE key pair with nil issuerKey")
			return
		}
		var size int
		issuerRSAKey, ok := s.issuerKey.(*rsa.PrivateKey)
		if ok {
			size = issuerRSAKey.N.BitLen()
		} else {
			size = 2048
		}
		eeKey, err = s.getRsaKey(size)
		if err != nil {
			err = errors.Wrapf(err, "failed to generate rsa private key of size %d", size)
			return
		}
		eePublicKey = eeKey.(*rsa.PrivateKey).Public()
	case cose.ES256, cose.ES384, cose.ES512:
		signer, err = cose.NewSigner(coseAlg, nil)
		if err != nil {
			err = errors.Wrapf(err, "failed to generate private key")
			return
		}
		eeKey = signer.PrivateKey
		eePublicKey = eeKey.(*ecdsa.PrivateKey).Public()
	}
	return
}

func expectHeadersAndGetKeyIDAndAlg(actual, expected *cose.Headers) (kidValue interface{}, alg *cose.Algorithm, err error) {
	if actual == nil || expected == nil {
		err = errors.New("xpi: cannot compare nil COSE headers")
		return
	}
	if len(actual.Unprotected) != len(expected.Unprotected) {
		err = fmt.Errorf("xpi: unexpected non-empty Unprotected headers got: %v", actual.Unprotected)
		return
	}
	if len(actual.Protected) != len(expected.Protected) {
		err = fmt.Errorf("xpi: unexpected Protected headers got: %v expected: %v", actual.Protected, expected.Protected)
		return
	}
	if _, ok := expected.Protected[algHeaderValue]; ok {
		algValue, ok := actual.Protected[algHeaderValue]
		if !ok {
			err = fmt.Errorf("xpi: missing expected alg in Protected Headers")
			return
		}
		if algInt, ok := algValue.(int); ok {
			alg = intToCOSEAlg(algInt)
		}
		if alg == nil {
			err = fmt.Errorf("xpi: alg %v is not supported", algValue)
			return
		}
	}
	if _, ok := expected.Protected[kidHeaderValue]; ok {
		kidValue, ok = actual.Protected[kidHeaderValue]
		if !ok {
			err = fmt.Errorf("xpi: missing expected kid in Protected Headers")
			return
		}
	}
	return
}

var (
	expectedMessageHeaders = &cose.Headers{
		Unprotected: map[interface{}]interface{}{},
		Protected: map[interface{}]interface{}{
			kidHeaderValue: nil,
		},
	}
	expectedSignatureHeaders = &cose.Headers{
		Unprotected: map[interface{}]interface{}{},
		Protected: map[interface{}]interface{}{
			kidHeaderValue: nil,
			algHeaderValue: nil,
		},
	}
)

// validateCOSESignatureStructureAndGetEECert checks whether a COSE
// signature structure is valid for an XPI and returns the parsed EE
// Cert from the protected header key id value. It does not verify the
// COSE signature bytes
func validateCOSESignatureStructureAndGetEECertAndAlg(sig *cose.Signature) (eeCert *x509.Certificate, algValue *cose.Algorithm, err error) {
	if sig == nil {
		err = errors.New("xpi: cannot validate nil COSE Signature")
		return
	}

	kidValue, algValue, err := expectHeadersAndGetKeyIDAndAlg(sig.Headers, expectedSignatureHeaders)
	if err != nil {
		err = errors.Wrapf(err, "xpi: got unexpected COSE Signature headers")
		return
	}

	kidBytes, ok := kidValue.([]byte)
	if !ok {
		err = fmt.Errorf("xpi: COSE Signature kid value is not a byte array")
		return
	}

	eeCert, err = x509.ParseCertificate(kidBytes)
	if err != nil {
		err = errors.Wrapf(err, "xpi: failed to parse X509 EE certificate from COSE Signature")
		return
	}

	return
}

// validateCOSEMessageStructureAndGetCerts checks whether a COSE
// SignMessage structure is valid for an XPI and returns the parsed
// intermediate and EE Certs from the protected header key id
// values. It does not verify the COSE signature bytes
func validateCOSEMessageStructureAndGetCertsAndAlgs(msg *cose.SignMessage) (intermediateCerts, eeCerts []*x509.Certificate, algs []*cose.Algorithm, err error) {
	if msg == nil {
		err = errors.New("xpi: cannot validate nil COSE SignMessage")
		return
	}
	if msg.Payload != nil {
		err = fmt.Errorf("xpi: expected SignMessage payload to be nil, but got %v", msg.Payload)
		return
	}
	kidValue, _, err := expectHeadersAndGetKeyIDAndAlg(msg.Headers, expectedMessageHeaders)
	if err != nil {
		err = errors.Wrapf(err, "xpi: got unexpected COSE SignMessage headers")
		return
	}

	// check that all kid values are bytes and decode into certs
	kidArray, ok := kidValue.([]interface{})
	if !ok {
		err = fmt.Errorf("xpi: expected SignMessage Protected Headers kid value to be an array got %v with type %T", kidValue, kidValue)
		return
	}
	for i, cert := range kidArray {
		certBytes, ok := cert.([]byte)
		if !ok {
			err = fmt.Errorf("xpi: expected SignMessage Protected Headers kid value %d to be a byte slice got %v with type %T", i, cert, cert)
			return
		}
		intermediateCert, parseErr := x509.ParseCertificate(certBytes)
		if parseErr != nil {
			err = errors.Wrapf(parseErr, "xpi: SignMessage Signature Protected Headers kid value %d does not decode to a parseable X509 cert", i)
			return
		}
		intermediateCerts = append(intermediateCerts, intermediateCert)
	}

	for i, sig := range msg.Signatures {
		eeCert, alg, sigErr := validateCOSESignatureStructureAndGetEECertAndAlg(&sig)
		if sigErr != nil {
			err = errors.Wrapf(sigErr, "xpi: cose signature %d is invalid", i)
			return
		}
		eeCerts = append(eeCerts, eeCert)
		algs = append(algs, alg)
	}

	return
}

// verifyCOSESignatures checks that:
//
// 1) COSE manifest and signature files are present
// 2) the PKCS7 manifest is present
// 3) the COSE and PKCS7 manifests do not include COSE files
// 4) we can decode the COSE signature and it has the right format for an XPI
// 5) the right number of signatures are present and all intermediate and end entity certs parse properly
// 6) **when a non-nil truststore is provided** that there is a trusted path from the included COSE EE certs to the signer cert using the provided intermediates
// 7) use the public keys from the EE certs to verify the COSE signature bytes
//
func verifyCOSESignatures(signedFile signer.SignedFile, truststore *x509.CertPool, signOptions Options) error {
	coseManifest, err := readFileFromZIP(signedFile, "META-INF/cose.manifest")
	if err != nil {
		return errors.Wrap(err, "xpi: failed to read META-INF/cose.manifest from signed zip")
	}
	coseMsgBytes, err := readFileFromZIP(signedFile, "META-INF/cose.sig")
	if err != nil {
		return errors.Wrap(err, "xpi: failed to read META-INF/cose.sig from signed zip")
	}
	pkcs7Manifest, err := readFileFromZIP(signedFile, "META-INF/manifest.mf")
	if err != nil {
		return errors.Wrap(err, "xpi: failed to read META-INF/manifest.mf from signed zip")
	}

	var coseFileNames = [][]byte{
		[]byte("Name: META-INF/cose.sig"),
		[]byte("Name: META-INF/cose.manifest"),
	}
	for _, coseFileName := range coseFileNames {
		if !bytes.Contains(pkcs7Manifest, coseFileName) {
			return fmt.Errorf("xpi: pkcs7 manifest does not contain the line: %s", coseFileName)
		}

		if bytes.Contains(coseManifest, coseFileName) {
			return fmt.Errorf("xpi: cose manifest contains the line: %s", coseFileName)
		}
	}

	xpiSig, unmarshalErr := Unmarshal(base64.StdEncoding.EncodeToString(coseMsgBytes), nil)
	if unmarshalErr != nil {
		return errors.Wrap(unmarshalErr, "xpi: error unmarshaling cose.sig")
	}
	if xpiSig != nil && xpiSig.signMessage != nil && len(xpiSig.signMessage.Signatures) != len(signOptions.COSEAlgorithms) {
		return fmt.Errorf("xpi: cose.sig contains %d signatures, but expected %d", len(xpiSig.signMessage.Signatures), len(signOptions.COSEAlgorithms))
	}

	intermediateCerts, eeCerts, algs, err := validateCOSEMessageStructureAndGetCertsAndAlgs(xpiSig.signMessage)
	if err != nil {
		return errors.Wrap(err, "xpi: cose.sig is not a valid COSE SignMessage")
	}

	// check that we can verify EE certs with the provided intermediates
	intermediates := x509.NewCertPool()
	for _, intermediateCert := range intermediateCerts {
		intermediates.AddCert(intermediateCert)
	}
	cndigest := sha256.Sum256([]byte(signOptions.ID))
	dnsName := fmt.Sprintf("%x.%x.addons.mozilla.org", cndigest[:16], cndigest[16:])

	var verifiers = []cose.Verifier{}

	for i, eeCert := range eeCerts {
		if signOptions.ID != eeCert.Subject.CommonName {
			return fmt.Errorf("xpi: EECert %d: id %s does not match cert cn %s", i, signOptions.ID, eeCert.Subject.CommonName)
		}
		opts := x509.VerifyOptions{
			DNSName:       dnsName,
			Roots:         truststore,
			Intermediates: intermediates,
		}
		if _, err := eeCert.Verify(opts); err != nil {
			return errors.Wrapf(err, "xpi: failed to verify EECert %d", i)
		}

		verifiers = append(verifiers, cose.Verifier{
			PublicKey: eeCert.PublicKey,
			Alg: algs[i],
		})
	}

	xpiSig.signMessage.Payload = coseManifest
	err = xpiSig.signMessage.Verify(nil, verifiers)
	if err != nil {
		return errors.Wrap(err, "xpi: failed to verify COSE SignMessage Signatures")
	}
	return nil
}

// issueCOSESignature returns a CBOR-marshalled COSE SignMessage
// after generating EE certs and signatures for the COSE algorithms
func (s *XPISigner) issueCOSESignature(cn string, manifest []byte, algs []*cose.Algorithm) (coseSig []byte, err error) {
	if s == nil {
		return nil, errors.New("xpi: cannot issue COSE Signature from nil XPISigner")
	}
	if s.issuerCert == nil {
		return nil, errors.New("xpi: cannot issue COSE Signature when XPISigner.issuerCert is nil")
	}
	if len(s.issuerCert.Raw) < 1 {
		return nil, errors.New("xpi: cannot issue COSE Signature when XPISigner.issuerCert is too short")
	}

	var (
		coseSigners []cose.Signer
		tmp         = cose.NewSignMessage()
		msg         = &tmp
	)
	msg.Payload = manifest

	// Add list of DER encoded intermediate certificates as message key id
	msg.Headers.Protected["kid"] = [][]byte{s.issuerCert.Raw[:]}

	for _, alg := range algs {
		// create a cert and key
		eeCert, eeKey, err := s.MakeEndEntity(cn, alg)
		if err != nil {
			return nil, err
		}

		// create a COSE.Signer
		signer, err := cose.NewSignerFromKey(alg, eeKey)
		if err != nil {
			return nil, errors.Wrap(err, "xpi: COSE signer creation failed")
		}
		coseSigners = append(coseSigners, *signer)

		// create a COSE Signature holder
		sig := cose.NewSignature()
		sig.Headers.Protected["alg"] = alg.Name
		sig.Headers.Protected["kid"] = eeCert.Raw[:]
		msg.AddSignature(sig)
	}

	// external_aad data must be nil and not byte("")
	err = msg.Sign(rand.Reader, nil, coseSigners)
	if err != nil {
		return nil, errors.Wrap(err, "xpi: COSE signing failed")
	}
	// for addons the signature is detached and the payload is always nil / null
	msg.Payload = nil

	coseSig, err = cose.Marshal(msg)
	if err != nil {
		return nil, errors.Wrap(err, "xpi: error serializing COSE signatures to CBOR")
	}

	return
}
<commit_msg>xpi: fix build error<commit_after>package xpi // import "go.mozilla.org/autograph/signer/xpi"

import (
	"bytes"
	"crypto"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"fmt"
	"strings"

	"github.com/pkg/errors"
	"go.mozilla.org/autograph/signer"
	"go.mozilla.org/cose"
)

const (
	// algHeaderValue compresses to 1 for the key "alg"
	algHeaderValue = 1
	// kidHeaderValue compresses to 4 for the key "kid"
	kidHeaderValue = 4
)

// stringToCOSEAlg returns the cose.Algorithm for a string or nil if
// the algorithm isn't implemented
func stringToCOSEAlg(s string) (v *cose.Algorithm) {
	switch strings.ToUpper(s) {
	case cose.PS256.Name:
		v = cose.PS256
	case cose.ES256.Name:
		v = cose.ES256
	case cose.ES384.Name:
		v = cose.ES384
	case cose.ES512.Name:
		v = cose.ES512
	default:
		v = nil
	}
	return v
}

// stringToCOSEAlg returns the cose.Algorithm for an int or nil if
// the algorithm isn't implemented
func intToCOSEAlg(i int) (v *cose.Algorithm) {
	switch i {
	case cose.PS256.Value:
		v = cose.PS256
	case cose.ES256.Value:
		v = cose.ES256
	case cose.ES384.Value:
		v = cose.ES384
	case cose.ES512.Value:
		v = cose.ES512
	default:
		v = nil
	}
	return v
}

// generateIssuerEEKeyPair returns a public and private key pair for
// the provided COSEAlgorithm
func (s *XPISigner) generateCOSEKeyPair(coseAlg *cose.Algorithm) (eeKey crypto.PrivateKey, eePublicKey crypto.PublicKey, err error) {
	var signer *cose.Signer

	switch coseAlg {
	case nil:
		err = errors.New("Cannot generate private key for nil cose Algorithm")
	case cose.PS256:
		if s.issuerKey == nil {
			err = errors.New("Cannot generate COSE key pair with nil issuerKey")
			return
		}
		var size int
		issuerRSAKey, ok := s.issuerKey.(*rsa.PrivateKey)
		if ok {
			size = issuerRSAKey.N.BitLen()
		} else {
			size = 2048
		}
		eeKey, err = s.getRsaKey(size)
		if err != nil {
			err = errors.Wrapf(err, "failed to generate rsa private key of size %d", size)
			return
		}
		eePublicKey = eeKey.(*rsa.PrivateKey).Public()
	case cose.ES256, cose.ES384, cose.ES512:
		signer, err = cose.NewSigner(coseAlg, nil)
		if err != nil {
			err = errors.Wrapf(err, "failed to generate private key")
			return
		}
		eeKey = signer.PrivateKey
		eePublicKey = eeKey.(*ecdsa.PrivateKey).Public()
	}
	return
}

func expectHeadersAndGetKeyIDAndAlg(actual, expected *cose.Headers) (kidValue interface{}, alg *cose.Algorithm, err error) {
	if actual == nil || expected == nil {
		err = errors.New("xpi: cannot compare nil COSE headers")
		return
	}
	if len(actual.Unprotected) != len(expected.Unprotected) {
		err = fmt.Errorf("xpi: unexpected non-empty Unprotected headers got: %v", actual.Unprotected)
		return
	}
	if len(actual.Protected) != len(expected.Protected) {
		err = fmt.Errorf("xpi: unexpected Protected headers got: %v expected: %v", actual.Protected, expected.Protected)
		return
	}
	if _, ok := expected.Protected[algHeaderValue]; ok {
		algValue, ok := actual.Protected[algHeaderValue]
		if !ok {
			err = fmt.Errorf("xpi: missing expected alg in Protected Headers")
			return
		}
		if algInt, ok := algValue.(int); ok {
			alg = intToCOSEAlg(algInt)
		}
		if alg == nil {
			err = fmt.Errorf("xpi: alg %v is not supported", algValue)
			return
		}
	}
	if _, ok := expected.Protected[kidHeaderValue]; ok {
		kidValue, ok = actual.Protected[kidHeaderValue]
		if !ok {
			err = fmt.Errorf("xpi: missing expected kid in Protected Headers")
			return
		}
	}
	return
}

var (
	expectedMessageHeaders = &cose.Headers{
		Unprotected: map[interface{}]interface{}{},
		Protected: map[interface{}]interface{}{
			kidHeaderValue: nil,
		},
	}
	expectedSignatureHeaders = &cose.Headers{
		Unprotected: map[interface{}]interface{}{},
		Protected: map[interface{}]interface{}{
			kidHeaderValue: nil,
			algHeaderValue: nil,
		},
	}
)

// validateCOSESignatureStructureAndGetEECert checks whether a COSE
// signature structure is valid for an XPI and returns the parsed EE
// Cert from the protected header key id value. It does not verify the
// COSE signature bytes
func validateCOSESignatureStructureAndGetEECertAndAlg(sig *cose.Signature) (eeCert *x509.Certificate, algValue *cose.Algorithm, err error) {
	if sig == nil {
		err = errors.New("xpi: cannot validate nil COSE Signature")
		return
	}

	kidValue, algValue, err := expectHeadersAndGetKeyIDAndAlg(sig.Headers, expectedSignatureHeaders)
	if err != nil {
		err = errors.Wrapf(err, "xpi: got unexpected COSE Signature headers")
		return
	}

	kidBytes, ok := kidValue.([]byte)
	if !ok {
		err = fmt.Errorf("xpi: COSE Signature kid value is not a byte array")
		return
	}

	eeCert, err = x509.ParseCertificate(kidBytes)
	if err != nil {
		err = errors.Wrapf(err, "xpi: failed to parse X509 EE certificate from COSE Signature")
		return
	}

	return
}

// validateCOSEMessageStructureAndGetCerts checks whether a COSE
// SignMessage structure is valid for an XPI and returns the parsed
// intermediate and EE Certs from the protected header key id
// values. It does not verify the COSE signature bytes
func validateCOSEMessageStructureAndGetCertsAndAlgs(msg *cose.SignMessage) (intermediateCerts, eeCerts []*x509.Certificate, algs []*cose.Algorithm, err error) {
	if msg == nil {
		err = errors.New("xpi: cannot validate nil COSE SignMessage")
		return
	}
	if msg.Payload != nil {
		err = fmt.Errorf("xpi: expected SignMessage payload to be nil, but got %v", msg.Payload)
		return
	}
	kidValue, _, err := expectHeadersAndGetKeyIDAndAlg(msg.Headers, expectedMessageHeaders)
	if err != nil {
		err = errors.Wrapf(err, "xpi: got unexpected COSE SignMessage headers")
		return
	}

	// check that all kid values are bytes and decode into certs
	kidArray, ok := kidValue.([]interface{})
	if !ok {
		err = fmt.Errorf("xpi: expected SignMessage Protected Headers kid value to be an array got %v with type %T", kidValue, kidValue)
		return
	}
	for i, cert := range kidArray {
		certBytes, ok := cert.([]byte)
		if !ok {
			err = fmt.Errorf("xpi: expected SignMessage Protected Headers kid value %d to be a byte slice got %v with type %T", i, cert, cert)
			return
		}
		intermediateCert, parseErr := x509.ParseCertificate(certBytes)
		if parseErr != nil {
			err = errors.Wrapf(parseErr, "xpi: SignMessage Signature Protected Headers kid value %d does not decode to a parseable X509 cert", i)
			return
		}
		intermediateCerts = append(intermediateCerts, intermediateCert)
	}

	for i, sig := range msg.Signatures {
		eeCert, alg, sigErr := validateCOSESignatureStructureAndGetEECertAndAlg(&sig)
		if sigErr != nil {
			err = errors.Wrapf(sigErr, "xpi: cose signature %d is invalid", i)
			return
		}
		eeCerts = append(eeCerts, eeCert)
		algs = append(algs, alg)
	}

	return
}

// verifyCOSESignatures checks that:
//
// 1) COSE manifest and signature files are present
// 2) the PKCS7 manifest is present
// 3) the COSE and PKCS7 manifests do not include COSE files
// 4) we can decode the COSE signature and it has the right format for an XPI
// 5) the right number of signatures are present and all intermediate and end entity certs parse properly
// 6) **when a non-nil truststore is provided** that there is a trusted path from the included COSE EE certs to the signer cert using the provided intermediates
// 7) use the public keys from the EE certs to verify the COSE signature bytes
//
func verifyCOSESignatures(signedFile signer.SignedFile, truststore *x509.CertPool, signOptions Options) error {
	coseManifest, err := readFileFromZIP(signedFile, "META-INF/cose.manifest")
	if err != nil {
		return errors.Wrap(err, "xpi: failed to read META-INF/cose.manifest from signed zip")
	}
	coseMsgBytes, err := readFileFromZIP(signedFile, "META-INF/cose.sig")
	if err != nil {
		return errors.Wrap(err, "xpi: failed to read META-INF/cose.sig from signed zip")
	}
	pkcs7Manifest, err := readFileFromZIP(signedFile, "META-INF/manifest.mf")
	if err != nil {
		return errors.Wrap(err, "xpi: failed to read META-INF/manifest.mf from signed zip")
	}

	var coseFileNames = [][]byte{
		[]byte("Name: META-INF/cose.sig"),
		[]byte("Name: META-INF/cose.manifest"),
	}
	for _, coseFileName := range coseFileNames {
		if !bytes.Contains(pkcs7Manifest, coseFileName) {
			return fmt.Errorf("xpi: pkcs7 manifest does not contain the line: %s", coseFileName)
		}

		if bytes.Contains(coseManifest, coseFileName) {
			return fmt.Errorf("xpi: cose manifest contains the line: %s", coseFileName)
		}
	}

	xpiSig, unmarshalErr := Unmarshal(base64.StdEncoding.EncodeToString(coseMsgBytes), nil)
	if unmarshalErr != nil {
		return errors.Wrap(unmarshalErr, "xpi: error unmarshaling cose.sig")
	}
	if xpiSig != nil && xpiSig.signMessage != nil && len(xpiSig.signMessage.Signatures) != len(signOptions.COSEAlgorithms) {
		return fmt.Errorf("xpi: cose.sig contains %d signatures, but expected %d", len(xpiSig.signMessage.Signatures), len(signOptions.COSEAlgorithms))
	}

	intermediateCerts, eeCerts, algs, err := validateCOSEMessageStructureAndGetCertsAndAlgs(xpiSig.signMessage)
	if err != nil {
		return errors.Wrap(err, "xpi: cose.sig is not a valid COSE SignMessage")
	}

	// check that we can verify EE certs with the provided intermediates
	intermediates := x509.NewCertPool()
	for _, intermediateCert := range intermediateCerts {
		intermediates.AddCert(intermediateCert)
	}
	cndigest := sha256.Sum256([]byte(signOptions.ID))
	dnsName := fmt.Sprintf("%x.%x.addons.mozilla.org", cndigest[:16], cndigest[16:])

	var verifiers = []cose.Verifier{}

	for i, eeCert := range eeCerts {
		if signOptions.ID != eeCert.Subject.CommonName {
			return fmt.Errorf("xpi: EECert %d: id %s does not match cert cn %s", i, signOptions.ID, eeCert.Subject.CommonName)
		}
		opts := x509.VerifyOptions{
			DNSName:       dnsName,
			Roots:         truststore,
			Intermediates: intermediates,
		}
		if _, err := eeCert.Verify(opts); err != nil {
			return errors.Wrapf(err, "xpi: failed to verify EECert %d", i)
		}

		verifiers = append(verifiers, cose.Verifier{
			PublicKey: eeCert.PublicKey,
			Alg: algs[i],
		})
	}

	xpiSig.signMessage.Payload = coseManifest
	err = xpiSig.signMessage.Verify(nil, verifiers)
	if err != nil {
		return errors.Wrap(err, "xpi: failed to verify COSE SignMessage Signatures")
	}
	return nil
}

// issueCOSESignature returns a CBOR-marshalled COSE SignMessage
// after generating EE certs and signatures for the COSE algorithms
func (s *XPISigner) issueCOSESignature(cn string, manifest []byte, algs []*cose.Algorithm) (coseSig []byte, err error) {
	if s == nil {
		return nil, errors.New("xpi: cannot issue COSE Signature from nil XPISigner")
	}
	if s.issuerCert == nil {
		return nil, errors.New("xpi: cannot issue COSE Signature when XPISigner.issuerCert is nil")
	}
	if len(s.issuerCert.Raw) < 1 {
		return nil, errors.New("xpi: cannot issue COSE Signature when XPISigner.issuerCert is too short")
	}

	var (
		coseSigners []cose.Signer
		msg         = cose.NewSignMessage()
	)
	msg.Payload = manifest

	// Add list of DER encoded intermediate certificates as message key id
	msg.Headers.Protected["kid"] = [][]byte{s.issuerCert.Raw[:]}

	for _, alg := range algs {
		// create a cert and key
		eeCert, eeKey, err := s.MakeEndEntity(cn, alg)
		if err != nil {
			return nil, err
		}

		// create a COSE.Signer
		signer, err := cose.NewSignerFromKey(alg, eeKey)
		if err != nil {
			return nil, errors.Wrap(err, "xpi: COSE signer creation failed")
		}
		coseSigners = append(coseSigners, *signer)

		// create a COSE Signature holder
		sig := cose.NewSignature()
		sig.Headers.Protected["alg"] = alg.Name
		sig.Headers.Protected["kid"] = eeCert.Raw[:]
		msg.AddSignature(sig)
	}

	// external_aad data must be nil and not byte("")
	err = msg.Sign(rand.Reader, nil, coseSigners)
	if err != nil {
		return nil, errors.Wrap(err, "xpi: COSE signing failed")
	}
	// for addons the signature is detached and the payload is always nil / null
	msg.Payload = nil

	coseSig, err = cose.Marshal(msg)
	if err != nil {
		return nil, errors.Wrap(err, "xpi: error serializing COSE signatures to CBOR")
	}

	return
}
<|endoftext|>