<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2018 Red Hat, Inc.
 *
 */

package network

import (
	"io/ioutil"
	"net"
	"os"
	"runtime"

	"github.com/golang/mock/gomock"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/vishvananda/netlink"

	v1 "kubevirt.io/client-go/api/v1"
	dutils "kubevirt.io/kubevirt/pkg/ephemeral-disk-utils"
	"kubevirt.io/kubevirt/pkg/network/cache"
	netdriver "kubevirt.io/kubevirt/pkg/network/driver"
	"kubevirt.io/kubevirt/pkg/virt-launcher/virtwrap/api"
)

var _ = Describe("Pod Network", func() {
	var mockNetwork *netdriver.MockNetworkHandler
	var ctrl *gomock.Controller
	var fakeMac net.HardwareAddr
	var tmpDir string
	var pid int
	var mtu int
	var cacheFactory cache.InterfaceCacheFactory

	BeforeEach(func() {
		dutils.MockDefaultOwnershipManager()
		var err error
		tmpDir, err = ioutil.TempDir("/tmp", "interface-cache")
		Expect(err).ToNot(HaveOccurred())
		cacheFactory = cache.NewInterfaceCacheFactoryWithBasePath(tmpDir)

		ctrl = gomock.NewController(GinkgoT())
		mockNetwork = netdriver.NewMockNetworkHandler(ctrl)
		mtu = 1410
		fakeMac, _ = net.ParseMAC("12:34:56:78:9A:BC")
		pid = os.Getpid()
	})

	AfterEach(func() {
		os.RemoveAll(tmpDir)
	})

	Context("on successful setup", func() {
		Context("Slirp Plug", func() {
			var (
				domain *api.Domain
				vmi    *v1.VirtualMachineInstance
			)

			BeforeEach(func() {
				domain = NewDomainWithSlirpInterface()
				api.NewDefaulter(runtime.GOARCH).SetObjectDefaults_Domain(domain)
				vmi = newVMISlirpInterface("testnamespace", "testVmName")
			})

			It("Should create an interface in the qemu command line and remove it from the interfaces", func() {
				specGenerator := newSlirpLibvirtSpecGenerator(&vmi.Spec.Domain.Devices.Interfaces[0], domain)
				Expect(specGenerator.generate()).To(Succeed())

				Expect(len(domain.Spec.Devices.Interfaces)).To(Equal(0))
				Expect(len(domain.Spec.QEMUCmd.QEMUArg)).To(Equal(2))
				Expect(domain.Spec.QEMUCmd.QEMUArg[0]).To(Equal(api.Arg{Value: "-device"}))
				Expect(domain.Spec.QEMUCmd.QEMUArg[1]).To(Equal(api.Arg{Value: "e1000,netdev=default,id=default"}))
			})

			It("Should append MAC address to qemu arguments if set", func() {
				vmi.Spec.Domain.Devices.Interfaces[0].MacAddress = "de-ad-00-00-be-af"
				specGenerator := newSlirpLibvirtSpecGenerator(&vmi.Spec.Domain.Devices.Interfaces[0], domain)
				Expect(specGenerator.generate()).To(Succeed())

				Expect(len(domain.Spec.Devices.Interfaces)).To(Equal(0))
				Expect(len(domain.Spec.QEMUCmd.QEMUArg)).To(Equal(2))
				Expect(domain.Spec.QEMUCmd.QEMUArg[0]).To(Equal(api.Arg{Value: "-device"}))
				Expect(domain.Spec.QEMUCmd.QEMUArg[1]).To(Equal(api.Arg{Value: "e1000,netdev=default,id=default,mac=de-ad-00-00-be-af"}))
			})
			It("Should create an interface in the qemu command line, remove it from the interfaces and leave the other interfaces inplace", func() {
				domain.Spec.Devices.Interfaces = append(domain.Spec.Devices.Interfaces, api.Interface{
					Model: &api.Model{
						Type: "virtio",
					},
					Type: "bridge",
					Source: api.InterfaceSource{
						Bridge: api.DefaultBridgeName,
					},
					Alias: api.NewUserDefinedAlias("default"),
				})
				specGenerator := newSlirpLibvirtSpecGenerator(&vmi.Spec.Domain.Devices.Interfaces[0], domain)
				Expect(specGenerator.generate()).To(Succeed())

				Expect(len(domain.Spec.Devices.Interfaces)).To(Equal(1))
				Expect(len(domain.Spec.QEMUCmd.QEMUArg)).To(Equal(2))
				Expect(domain.Spec.QEMUCmd.QEMUArg[0]).To(Equal(api.Arg{Value: "-device"}))
				Expect(domain.Spec.QEMUCmd.QEMUArg[1]).To(Equal(api.Arg{Value: "e1000,netdev=default,id=default"}))
			})
		})
		Context("Macvtap plug", func() {
			var podnic *podNIC

			BeforeEach(func() {
				vmi := newVMIMacvtapInterface("testnamespace", "testVmName", "default")
				var err error
				podnic, err = newPhase1PodNIC(vmi, &vmi.Spec.Networks[0], mockNetwork, cacheFactory, &pid)
				Expect(err).ToNot(HaveOccurred())
				macvtapInterface := &netlink.GenericLink{LinkAttrs: netlink.LinkAttrs{Name: podnic.podInterfaceName, MTU: mtu, HardwareAddr: fakeMac}}
				mockNetwork.EXPECT().LinkByName(podnic.podInterfaceName).Return(macvtapInterface, nil)
			})

			It("Should pass a non-privileged macvtap interface to qemu", func() {
				domain := NewDomainWithMacvtapInterface("default")

				api.NewDefaulter(runtime.GOARCH).SetObjectDefaults_Domain(domain)
				specGenerator := podnic.newLibvirtSpecGenerator(domain)

				Expect(specGenerator.generate()).To(Succeed())

				Expect(len(domain.Spec.Devices.Interfaces)).To(Equal(1), "should have a single interface")
				Expect(domain.Spec.Devices.Interfaces[0].Target).To(Equal(&api.InterfaceTarget{Device: podnic.podInterfaceName, Managed: "no"}), "should have an unmanaged interface")
				Expect(domain.Spec.Devices.Interfaces[0].MAC).To(Equal(&api.MAC{MAC: fakeMac.String()}), "should have the expected MAC address")
				Expect(domain.Spec.Devices.Interfaces[0].MTU).To(Equal(&api.MTU{Size: "1410"}), "should have the expected MTU")
			})
		})
	})
})
<commit_msg>network, macvtap: do not rely on podnic to unit test spec generators<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2018 Red Hat, Inc.
 *
 */

package network

import (
	"io/ioutil"
	"net"
	"os"
	"runtime"

	"github.com/golang/mock/gomock"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/vishvananda/netlink"

	v1 "kubevirt.io/client-go/api/v1"
	dutils "kubevirt.io/kubevirt/pkg/ephemeral-disk-utils"
	netdriver "kubevirt.io/kubevirt/pkg/network/driver"
	"kubevirt.io/kubevirt/pkg/virt-launcher/virtwrap/api"
)

var _ = Describe("Pod Network", func() {
	var mockNetwork *netdriver.MockNetworkHandler
	var ctrl *gomock.Controller
	var fakeMac net.HardwareAddr
	var tmpDir string
	var mtu int

	BeforeEach(func() {
		dutils.MockDefaultOwnershipManager()
		var err error
		tmpDir, err = ioutil.TempDir("/tmp", "interface-cache")
		Expect(err).ToNot(HaveOccurred())

		ctrl = gomock.NewController(GinkgoT())
		mockNetwork = netdriver.NewMockNetworkHandler(ctrl)
		mtu = 1410
		fakeMac, _ = net.ParseMAC("12:34:56:78:9A:BC")
	})

	AfterEach(func() {
		os.RemoveAll(tmpDir)
	})

	Context("on successful setup", func() {
		Context("Slirp Plug", func() {
			var (
				domain *api.Domain
				vmi    *v1.VirtualMachineInstance
			)

			BeforeEach(func() {
				domain = NewDomainWithSlirpInterface()
				api.NewDefaulter(runtime.GOARCH).SetObjectDefaults_Domain(domain)
				vmi = newVMISlirpInterface("testnamespace", "testVmName")
			})

			It("Should create an interface in the qemu command line and remove it from the interfaces", func() {
				specGenerator := newSlirpLibvirtSpecGenerator(&vmi.Spec.Domain.Devices.Interfaces[0], domain)
				Expect(specGenerator.generate()).To(Succeed())

				Expect(len(domain.Spec.Devices.Interfaces)).To(Equal(0))
				Expect(len(domain.Spec.QEMUCmd.QEMUArg)).To(Equal(2))
				Expect(domain.Spec.QEMUCmd.QEMUArg[0]).To(Equal(api.Arg{Value: "-device"}))
				Expect(domain.Spec.QEMUCmd.QEMUArg[1]).To(Equal(api.Arg{Value: "e1000,netdev=default,id=default"}))
			})

			It("Should append MAC address to qemu arguments if set", func() {
				vmi.Spec.Domain.Devices.Interfaces[0].MacAddress = "de-ad-00-00-be-af"
				specGenerator := newSlirpLibvirtSpecGenerator(&vmi.Spec.Domain.Devices.Interfaces[0], domain)
				Expect(specGenerator.generate()).To(Succeed())

				Expect(len(domain.Spec.Devices.Interfaces)).To(Equal(0))
				Expect(len(domain.Spec.QEMUCmd.QEMUArg)).To(Equal(2))
				Expect(domain.Spec.QEMUCmd.QEMUArg[0]).To(Equal(api.Arg{Value: "-device"}))
				Expect(domain.Spec.QEMUCmd.QEMUArg[1]).To(Equal(api.Arg{Value: "e1000,netdev=default,id=default,mac=de-ad-00-00-be-af"}))
			})
			It("Should create an interface in the qemu command line, remove it from the interfaces and leave the other interfaces inplace", func() {
				domain.Spec.Devices.Interfaces = append(domain.Spec.Devices.Interfaces, api.Interface{
					Model: &api.Model{
						Type: "virtio",
					},
					Type: "bridge",
					Source: api.InterfaceSource{
						Bridge: api.DefaultBridgeName,
					},
					Alias: api.NewUserDefinedAlias("default"),
				})
				specGenerator := newSlirpLibvirtSpecGenerator(&vmi.Spec.Domain.Devices.Interfaces[0], domain)
				Expect(specGenerator.generate()).To(Succeed())

				Expect(len(domain.Spec.Devices.Interfaces)).To(Equal(1))
				Expect(len(domain.Spec.QEMUCmd.QEMUArg)).To(Equal(2))
				Expect(domain.Spec.QEMUCmd.QEMUArg[0]).To(Equal(api.Arg{Value: "-device"}))
				Expect(domain.Spec.QEMUCmd.QEMUArg[1]).To(Equal(api.Arg{Value: "e1000,netdev=default,id=default"}))
			})
		})
		Context("Macvtap plug", func() {
			const primaryPodIfaceName = "eth0"

			var (
				domain        *api.Domain
				specGenerator *MacvtapLibvirtSpecGenerator
			)

			BeforeEach(func() {
				domain = NewDomainWithMacvtapInterface("default")
				api.NewDefaulter(runtime.GOARCH).SetObjectDefaults_Domain(domain)
				vmi := newVMIMacvtapInterface("testnamespace", "testVmName", "default")
				macvtapInterface := &netlink.GenericLink{LinkAttrs: netlink.LinkAttrs{Name: primaryPodIfaceName, MTU: mtu, HardwareAddr: fakeMac}}
				mockNetwork.EXPECT().LinkByName(primaryPodIfaceName).Return(macvtapInterface, nil)
				specGenerator = newMacvtapLibvirtSpecGenerator(
					&vmi.Spec.Domain.Devices.Interfaces[0], domain, primaryPodIfaceName, mockNetwork)
			})

			It("Should pass a non-privileged macvtap interface to qemu", func() {
				Expect(specGenerator.generate()).To(Succeed())

				Expect(len(domain.Spec.Devices.Interfaces)).To(Equal(1), "should have a single interface")
				Expect(domain.Spec.Devices.Interfaces[0].Target).To(Equal(&api.InterfaceTarget{Device: primaryPodIfaceName, Managed: "no"}), "should have an unmanaged interface")
				Expect(domain.Spec.Devices.Interfaces[0].MAC).To(Equal(&api.MAC{MAC: fakeMac.String()}), "should have the expected MAC address")
				Expect(domain.Spec.Devices.Interfaces[0].MTU).To(Equal(&api.MTU{Size: "1410"}), "should have the expected MTU")
			})
		})
	})
})
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"strconv"
	"strings"
	"text/tabwriter"

	"github.com/bradfitz/slice"
)

// Character is the type representing a role playing character
type Character struct {
	Name            string
	Backgrounds     map[string]Background
	Aptitudes       map[string]Aptitude
	Characteristics map[string]Characteristic
	Skills          map[string]Skill
	Talents         map[string]Talent
	Gauges          map[string]Gauge
	Rules           map[string]Rule
	Experience      int
	Spent           int
}

// NewCharacter creates a new character from the given sheet and universe.
func NewCharacter(universe Universe, sheet Sheet) (Character, error) {

	// Create a character
	c := Character{
		Name:            sheet.Header.Name,
		Backgrounds:     make(map[string]Background),
		Aptitudes:       make(map[string]Aptitude),
		Characteristics: make(map[string]Characteristic),
		Skills:          make(map[string]Skill),
		Talents:         make(map[string]Talent),
		Gauges:          make(map[string]Gauge),
		Rules:           make(map[string]Rule),
		Experience:      0,
		Spent:           0,
	}

	// The characteristics described in the header of the sheet are parsed as upgrades
	for _, upgrade := range sheet.Characteristics {

		// Get the characteristic from the universe
		characteristic, found := universe.FindCharacteristic(upgrade.Name)
		if !found {
			return c, NewError(UndefinedCharacteristic, upgrade.Line)
		}

		// Check it is not already applied
		_, found = c.Characteristics[characteristic.Name]
		if found {
			return c, NewError(DuplicateCharacteristic, upgrade.Line)
		}

		// Apply the upgrade
		err := c.ApplyCharacteristicUpgrade(characteristic, upgrade)
		if err != nil {
			return c, err
		}
	}

	// Next are the backgrounds
	for typ, metas := range sheet.Header.Metas {

		for _, meta := range metas {

			// Find the background corresponding to the meta
			background, err := universe.FindBackground(typ, meta.Label)
			if err != nil {
				return c, err
			}

			err = c.ApplyBackground(background, universe)
			if err != nil {
				return c, err
			}
		}
	}

	// Next are the sessions
	for _, session := range sheet.Sessions {

		// Apply the experience gain if needed
		if session.Reward != nil {
			c.Experience += *session.Reward
		}

		// Apply each upgrade in order
		for _, upgrade := range session.Upgrades {
			err := c.ApplyUpgrade(upgrade, universe)
			if err != nil {
				return c, err
			}
		}
	}

	return c, nil
}

// CountMatchingAptitudes return the number of aptitudes of the given slice
// that are in the character's aptitudes.
func (character Character) CountMatchingAptitudes(aptitudes []Aptitude) int {

	count := 0
	for _, aptitude := range aptitudes {
		if _, found := character.Aptitudes[string(aptitude)]; found {
			count++
		}
	}
	return count
}

// ApplyBackground changes the character's trait according to the history values
func (character *Character) ApplyBackground(background Background, universe Universe) error {

	cost := 0

	// For each upgrade associated to the history, apply each option.
	for _, upgrade := range background.Upgrades {
		err := character.ApplyUpgrade(Upgrade{
			Mark: MarkSpecial,
			Name: upgrade,
			Cost: &cost,
		}, universe)
		if err != nil {
			return err
		}
	}

	// Add the background to the character's backgrounds
	character.Backgrounds[background.Name] = background

	return nil
}

// ApplyUpgrade changes the character's attributes according to the given upgrade.
func (character *Character) ApplyUpgrade(upgrade Upgrade, universe Universe) error {

	var err error

	// Find the attribute corresponding to the upgrade., and initialize a new
	// rule if there isn't any.
	coster, found := universe.FindCoster(upgrade.Name)
	if !found {
		coster = Rule{
			Name: upgrade.Name,
		}
	}

	// If no cost is defined, compute it on the fly.
	if upgrade.Cost == nil {
		cost, err := coster.Cost(universe, *character)
		if err != nil {
			return err
		}

		upgrade.Cost = &cost
	}

	// Update the spent experience.
	character.Spent += *upgrade.Cost

	// Apply the upgrade depending on the target attribute.
	switch attribute := coster.(type) {
	case Characteristic:
		err = character.ApplyCharacteristicUpgrade(attribute, upgrade)
	case Skill:
		err = character.ApplySkillUpgrade(attribute, upgrade)
	case Talent:
		err = character.ApplyTalentUpgrade(attribute, upgrade)
	case Aptitude:
		err = character.ApplyAptitudeUpgrade(attribute, upgrade)
	case Gauge:
		err = character.ApplyGaugeUpgrade(attribute, upgrade)
	case Rule:
		err = character.ApplyRuleUpgrade(attribute, upgrade)
	}

	return err
}

func (character *Character) ApplyCharacteristicUpgrade(characteristic Characteristic, upgrade Upgrade) error {

	// Get the attribute from the character's characteristic map.
	c, found := character.Characteristics[characteristic.Name]
	if !found {
		c = characteristic
	}

	// Increment the tier if the mark is default.
	if upgrade.Mark == MarkDefault {
		c.Tier++
	}

	// Parse the characteristic's upgrade value.
	raw := strings.TrimSpace(strings.TrimLeft(upgrade.Name, characteristic.Name))
	value, err := strconv.Atoi(raw)
	if err != nil {
		return NewError(InvalidCharacteristicValue)
	}

	// Update the characteristic value.
	if strings.HasPrefix(raw, "+") || strings.HasPrefix(raw, "-") {
		c.Value += value
	} else {
		c.Value = value
	}

	character.Characteristics[c.Name] = c

	return nil
}

func (character *Character) ApplySkillUpgrade(skill Skill, upgrade Upgrade) error {

	// Get the skill from the character's skill map.
	s, found := character.Skills[skill.FullName()]
	if !found {
		s = skill
	}

	// Increment the tier if the mark is default.
	if upgrade.Mark == MarkDefault {
		s.Tier++
	}

	// Put the skill back on the map.
	character.Skills[skill.FullName()] = s

	return nil
}

func (character *Character) ApplyTalentUpgrade(talent Talent, upgrade Upgrade) error {

	// Get the talent from the character.
	t, found := character.Talents[talent.FullName()]
	if !found {
		t = talent
	}

	// Increment the value of the talent.
	t.Value++

	// Put it back on the map.
	character.Talents[talent.FullName()] = t

	return nil
}

func (character *Character) ApplyAptitudeUpgrade(aptitude Aptitude, upgrade Upgrade) error {

	// Add the aptitude to the character's aptitudes.
	character.Aptitudes[string(aptitude)] = aptitude

	return nil
}

func (character *Character) ApplyGaugeUpgrade(gauge Gauge, upgrade Upgrade) error {

	// Get the gauge from the character.
	g, found := character.Gauges[gauge.Name]
	if !found {
		g = gauge
	}

	// Parse the gauge's upgrade value.
	raw := strings.TrimSpace(strings.TrimLeft(upgrade.Name, g.Name))
	value, err := strconv.Atoi(raw)
	if err != nil {
		return NewError(InvalidGaugeValue)
	}

	// Update the gauge value.
	if strings.HasPrefix(raw, "+") || strings.HasPrefix(raw, "-") {
		g.Value += value
	} else {
		g.Value = value
	}

	// Set the gauge back on the map.
	character.Gauges[g.Name] = g

	return nil
}

func (character *Character) ApplyRuleUpgrade(rule Rule, upgrade Upgrade) error {

	// Add the rule to the character's rules.
	character.Rules[rule.Name] = rule

	return nil
}

// Print the character sheet on the screen
func (character Character) Print() {
	// Print the name
	fmt.Printf("%s\t%s\n", theme.Title("Name"), character.Name)

	// Print the backgrounds
	backgrounds := []Background{}

	for _, background := range character.Backgrounds {
		backgrounds = append(backgrounds, background)
	}

	slice.Sort(backgrounds, func(i, j int) bool {
		if backgrounds[i].Type != backgrounds[j].Type {
			return backgrounds[i].Type < backgrounds[j].Type
		}

		return backgrounds[i].Name < backgrounds[j].Name
	})

	for _, background := range backgrounds {
		fmt.Printf("%s\t%s\n", theme.Title(strings.Title(background.Type)), strings.Title(background.Name))
	}

	// Print the experience
	fmt.Printf("\n%s\t%d/%d\n", theme.Title("Experience"), character.Spent, character.Experience)

	// Print the characteristics
	fmt.Printf("\n%s\n", theme.Title("Characteristics"))

	characteristics := []Characteristic{}
	for _, characteristic := range character.Characteristics {
		characteristics = append(characteristics, characteristic)
	}

	slice.Sort(characteristics, func(i, j int) bool {
		return characteristics[i].Name < characteristics[j].Name
	})

	for _, characteristic := range characteristics {
		fmt.Printf("%s\t%s\n", characteristic.Name, theme.Value(characteristic.Value))
	}

	// Print the gauges
	fmt.Printf("\n%s\n", theme.Title("Gauges"))

	gauges := []Gauge{}
	for _, gauge := range character.Gauges {
		gauges = append(gauges, gauge)
	}

	slice.Sort(gauges, func(i, j int) bool {
		return gauges[i].Name < gauges[j].Name
	})

	for _, gauge := range gauges {
		fmt.Printf("%s\t%s\n", gauge.Name, theme.Value(gauge.Value))
	}

	// Print the skills using a tabwriter
	fmt.Printf("\n%s\n", theme.Title("Skills"))

	skills := []Skill{}
	for _, skill := range character.Skills {
		skills = append(skills, skill)
	}

	slice.Sort(skills, func(i, j int) bool {
		return skills[i].FullName() < skills[j].FullName()
	})

	w := tabwriter.NewWriter(os.Stdout, 10, 1, 2, ' ', 0)
	for _, skill := range skills {
		fmt.Fprintf(w, "%s\t+%s\n", strings.Title(skill.FullName()), theme.Value(skill.Tier*10))
	}
	w.Flush()

	// Print the talents
	fmt.Printf("\n%s\n", theme.Title("Talents"))

	talents := []Talent{}
	for _, talent := range character.Talents {
		talents = append(talents, talent)
	}

	slice.Sort(talents, func(i, j int) bool {
		return talents[i].FullName() < talents[j].FullName()
	})

	for _, talent := range talents {
		if talent.Value != 1 {
			fmt.Printf("%s (%s)\n", strings.Title(talent.FullName()), theme.Value(talent.Value))
		} else {
			fmt.Printf("%s\n", strings.Title(talent.FullName()))
		}
	}

	// Print the special rules
	fmt.Printf("\n%s\n", theme.Title("Rules"))

	rules := []Rule{}

	for _, rule := range character.Rules {
		rules = append(rules, rule)
	}

	slice.Sort(rules, func(i, j int) bool {
		return rules[i].Name < rules[j].Name
	})

	for _, rule := range rules {
		fmt.Printf("%s\t%s\n", strings.Title(rule.Name), rule.Description)
	}
}
<commit_msg>Prints characters aptitudes<commit_after>package main

import (
	"fmt"
	"os"
	"strconv"
	"strings"
	"text/tabwriter"

	"github.com/bradfitz/slice"
)

// Character is the type representing a role playing character
type Character struct {
	Name            string
	Backgrounds     map[string]Background
	Aptitudes       map[string]Aptitude
	Characteristics map[string]Characteristic
	Skills          map[string]Skill
	Talents         map[string]Talent
	Gauges          map[string]Gauge
	Rules           map[string]Rule
	Experience      int
	Spent           int
}

// NewCharacter creates a new character from the given sheet and universe.
func NewCharacter(universe Universe, sheet Sheet) (Character, error) {

	// Create a character
	c := Character{
		Name:            sheet.Header.Name,
		Backgrounds:     make(map[string]Background),
		Aptitudes:       make(map[string]Aptitude),
		Characteristics: make(map[string]Characteristic),
		Skills:          make(map[string]Skill),
		Talents:         make(map[string]Talent),
		Gauges:          make(map[string]Gauge),
		Rules:           make(map[string]Rule),
		Experience:      0,
		Spent:           0,
	}

	// The characteristics described in the header of the sheet are parsed as upgrades
	for _, upgrade := range sheet.Characteristics {

		// Get the characteristic from the universe
		characteristic, found := universe.FindCharacteristic(upgrade.Name)
		if !found {
			return c, NewError(UndefinedCharacteristic, upgrade.Line)
		}

		// Check it is not already applied
		_, found = c.Characteristics[characteristic.Name]
		if found {
			return c, NewError(DuplicateCharacteristic, upgrade.Line)
		}

		// Apply the upgrade
		err := c.ApplyCharacteristicUpgrade(characteristic, upgrade)
		if err != nil {
			return c, err
		}
	}

	// Next are the backgrounds
	for typ, metas := range sheet.Header.Metas {

		for _, meta := range metas {

			// Find the background corresponding to the meta
			background, err := universe.FindBackground(typ, meta.Label)
			if err != nil {
				return c, err
			}

			err = c.ApplyBackground(background, universe)
			if err != nil {
				return c, err
			}
		}
	}

	// Next are the sessions
	for _, session := range sheet.Sessions {

		// Apply the experience gain if needed
		if session.Reward != nil {
			c.Experience += *session.Reward
		}

		// Apply each upgrade in order
		for _, upgrade := range session.Upgrades {
			err := c.ApplyUpgrade(upgrade, universe)
			if err != nil {
				return c, err
			}
		}
	}

	return c, nil
}

// CountMatchingAptitudes return the number of aptitudes of the given slice
// that are in the character's aptitudes.
func (character Character) CountMatchingAptitudes(aptitudes []Aptitude) int {

	count := 0
	for _, aptitude := range aptitudes {
		if _, found := character.Aptitudes[string(aptitude)]; found {
			count++
		}
	}
	return count
}

// ApplyBackground changes the character's trait according to the history values
func (character *Character) ApplyBackground(background Background, universe Universe) error {

	cost := 0

	// For each upgrade associated to the history, apply each option.
	for _, upgrade := range background.Upgrades {
		err := character.ApplyUpgrade(Upgrade{
			Mark: MarkSpecial,
			Name: upgrade,
			Cost: &cost,
		}, universe)
		if err != nil {
			return err
		}
	}

	// Add the background to the character's backgrounds
	character.Backgrounds[background.Name] = background

	return nil
}

// ApplyUpgrade changes the character's attributes according to the given upgrade.
func (character *Character) ApplyUpgrade(upgrade Upgrade, universe Universe) error {

	var err error

	// Find the attribute corresponding to the upgrade., and initialize a new
	// rule if there isn't any.
	coster, found := universe.FindCoster(upgrade.Name)
	if !found {
		coster = Rule{
			Name: upgrade.Name,
		}
	}

	// If no cost is defined, compute it on the fly.
	if upgrade.Cost == nil {
		cost, err := coster.Cost(universe, *character)
		if err != nil {
			return err
		}

		upgrade.Cost = &cost
	}

	// Update the spent experience.
	character.Spent += *upgrade.Cost

	// Apply the upgrade depending on the target attribute.
	switch attribute := coster.(type) {
	case Characteristic:
		err = character.ApplyCharacteristicUpgrade(attribute, upgrade)
	case Skill:
		err = character.ApplySkillUpgrade(attribute, upgrade)
	case Talent:
		err = character.ApplyTalentUpgrade(attribute, upgrade)
	case Aptitude:
		err = character.ApplyAptitudeUpgrade(attribute, upgrade)
	case Gauge:
		err = character.ApplyGaugeUpgrade(attribute, upgrade)
	case Rule:
		err = character.ApplyRuleUpgrade(attribute, upgrade)
	}

	return err
}

func (character *Character) ApplyCharacteristicUpgrade(characteristic Characteristic, upgrade Upgrade) error {

	// Get the attribute from the character's characteristic map.
	c, found := character.Characteristics[characteristic.Name]
	if !found {
		c = characteristic
	}

	// Increment the tier if the mark is default.
	if upgrade.Mark == MarkDefault {
		c.Tier++
	}

	// Parse the characteristic's upgrade value.
	raw := strings.TrimSpace(strings.TrimLeft(upgrade.Name, characteristic.Name))
	value, err := strconv.Atoi(raw)
	if err != nil {
		return NewError(InvalidCharacteristicValue)
	}

	// Update the characteristic value.
	if strings.HasPrefix(raw, "+") || strings.HasPrefix(raw, "-") {
		c.Value += value
	} else {
		c.Value = value
	}

	character.Characteristics[c.Name] = c

	return nil
}

func (character *Character) ApplySkillUpgrade(skill Skill, upgrade Upgrade) error {

	// Get the skill from the character's skill map.
	s, found := character.Skills[skill.FullName()]
	if !found {
		s = skill
	}

	// Increment the tier if the mark is default.
	if upgrade.Mark == MarkDefault {
		s.Tier++
	}

	// Put the skill back on the map.
	character.Skills[skill.FullName()] = s

	return nil
}

func (character *Character) ApplyTalentUpgrade(talent Talent, upgrade Upgrade) error {

	// Get the talent from the character.
	t, found := character.Talents[talent.FullName()]
	if !found {
		t = talent
	}

	// Increment the value of the talent.
	t.Value++

	// Put it back on the map.
	character.Talents[talent.FullName()] = t

	return nil
}

func (character *Character) ApplyAptitudeUpgrade(aptitude Aptitude, upgrade Upgrade) error {

	// Add the aptitude to the character's aptitudes.
	character.Aptitudes[string(aptitude)] = aptitude

	return nil
}

func (character *Character) ApplyGaugeUpgrade(gauge Gauge, upgrade Upgrade) error {

	// Get the gauge from the character.
	g, found := character.Gauges[gauge.Name]
	if !found {
		g = gauge
	}

	// Parse the gauge's upgrade value.
	raw := strings.TrimSpace(strings.TrimLeft(upgrade.Name, g.Name))
	value, err := strconv.Atoi(raw)
	if err != nil {
		return NewError(InvalidGaugeValue)
	}

	// Update the gauge value.
	if strings.HasPrefix(raw, "+") || strings.HasPrefix(raw, "-") {
		g.Value += value
	} else {
		g.Value = value
	}

	// Set the gauge back on the map.
	character.Gauges[g.Name] = g

	return nil
}

func (character *Character) ApplyRuleUpgrade(rule Rule, upgrade Upgrade) error {

	// Add the rule to the character's rules.
	character.Rules[rule.Name] = rule

	return nil
}

// Print the character sheet on the screen
func (character Character) Print() {
	// Print the name
	fmt.Printf("%s\t%s\n", theme.Title("Name"), character.Name)

	// Print the backgrounds
	backgrounds := []Background{}

	for _, background := range character.Backgrounds {
		backgrounds = append(backgrounds, background)
	}

	slice.Sort(backgrounds, func(i, j int) bool {
		if backgrounds[i].Type != backgrounds[j].Type {
			return backgrounds[i].Type < backgrounds[j].Type
		}

		return backgrounds[i].Name < backgrounds[j].Name
	})

	for _, background := range backgrounds {
		fmt.Printf("%s\t%s\n", theme.Title(strings.Title(background.Type)), strings.Title(background.Name))
	}

	// Print the aptitudes
	aptitudes := []Aptitude{}

	for _, aptitude := range character.Aptitudes {
		aptitudes = append(aptitudes, aptitude)
	}

	slice.Sort(aptitudes, func(i, j int) bool {
		return aptitudes[i] < aptitudes[j]
	})

	fmt.Printf("\n%s\n", theme.Title("Aptitudes"))
	for _, aptitude := range aptitudes {
		fmt.Printf("\t%s\n", strings.Title(string(aptitude)))
	}

	// Print the experience
	fmt.Printf("\n%s\t%d/%d\n", theme.Title("Experience"), character.Spent, character.Experience)

	// Print the characteristics
	fmt.Printf("\n%s\n", theme.Title("Characteristics"))

	characteristics := []Characteristic{}
	for _, characteristic := range character.Characteristics {
		characteristics = append(characteristics, characteristic)
	}

	slice.Sort(characteristics, func(i, j int) bool {
		return characteristics[i].Name < characteristics[j].Name
	})

	for _, characteristic := range characteristics {
		fmt.Printf("%s\t%s\n", characteristic.Name, theme.Value(characteristic.Value))
	}

	// Print the gauges
	fmt.Printf("\n%s\n", theme.Title("Gauges"))

	gauges := []Gauge{}
	for _, gauge := range character.Gauges {
		gauges = append(gauges, gauge)
	}

	slice.Sort(gauges, func(i, j int) bool {
		return gauges[i].Name < gauges[j].Name
	})

	for _, gauge := range gauges {
		fmt.Printf("%s\t%s\n", gauge.Name, theme.Value(gauge.Value))
	}

	// Print the skills using a tabwriter
	fmt.Printf("\n%s\n", theme.Title("Skills"))

	skills := []Skill{}
	for _, skill := range character.Skills {
		skills = append(skills, skill)
	}

	slice.Sort(skills, func(i, j int) bool {
		return skills[i].FullName() < skills[j].FullName()
	})

	w := tabwriter.NewWriter(os.Stdout, 10, 1, 2, ' ', 0)
	for _, skill := range skills {
		fmt.Fprintf(w, "%s\t+%s\n", strings.Title(skill.FullName()), theme.Value(skill.Tier*10))
	}
	w.Flush()

	// Print the talents
	fmt.Printf("\n%s\n", theme.Title("Talents"))

	talents := []Talent{}
	for _, talent := range character.Talents {
		talents = append(talents, talent)
	}

	slice.Sort(talents, func(i, j int) bool {
		return talents[i].FullName() < talents[j].FullName()
	})

	for _, talent := range talents {
		if talent.Value != 1 {
			fmt.Printf("%s (%s)\n", strings.Title(talent.FullName()), theme.Value(talent.Value))
		} else {
			fmt.Printf("%s\n", strings.Title(talent.FullName()))
		}
	}

	// Print the special rules
	fmt.Printf("\n%s\n", theme.Title("Rules"))

	rules := []Rule{}

	for _, rule := range character.Rules {
		rules = append(rules, rule)
	}

	slice.Sort(rules, func(i, j int) bool {
		return rules[i].Name < rules[j].Name
	})

	for _, rule := range rules {
		fmt.Printf("%s\t%s\n", strings.Title(rule.Name), rule.Description)
	}
}
<|endoftext|>


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

package linux

import (
	"gvisor.googlesource.com/gvisor/pkg/abi/linux"
	"gvisor.googlesource.com/gvisor/pkg/sentry/arch"
	"gvisor.googlesource.com/gvisor/pkg/sentry/fs"
	"gvisor.googlesource.com/gvisor/pkg/sentry/kernel"
	"gvisor.googlesource.com/gvisor/pkg/sentry/usermem"
	"gvisor.googlesource.com/gvisor/pkg/syserror"
)

// Mount implements Linux syscall mount(2).
func Mount(t *kernel.Task, args arch.SyscallArguments) (uintptr, *kernel.SyscallControl, error) {
	sourceAddr := args[0].Pointer()
	targetAddr := args[1].Pointer()
	typeAddr := args[2].Pointer()
	flags := args[3].Uint64()
	dataAddr := args[4].Pointer()

	fsType, err := t.CopyInString(typeAddr, usermem.PageSize)
	if err != nil {
		return 0, nil, err
	}

	sourcePath, _, err := copyInPath(t, sourceAddr, true /* allowEmpty */)
	if err != nil {
		return 0, nil, err
	}

	targetPath, _, err := copyInPath(t, targetAddr, false /* allowEmpty */)
	if err != nil {
		return 0, nil, err
	}

	// In Linux, a full page is always copied in regardless of null
	// character placement, and the address is passed to each file system.
	// Most file systems always treat this data as a string, though, and so
	// do all of the ones we implement.
	data, err := t.CopyInString(dataAddr, usermem.PageSize)
	if err != nil {
		return 0, nil, err
	}

	// Ignore magic value that was required before Linux 2.4.
	if flags&linux.MS_MGC_MSK == linux.MS_MGC_VAL {
		flags = flags &^ linux.MS_MGC_MSK
	}

	// Must have CAP_SYS_ADMIN in the mount namespace's associated user
	// namespace.
	if !t.HasCapabilityIn(linux.CAP_SYS_ADMIN, t.MountNamespace().UserNamespace()) {
		return 0, nil, syserror.EPERM
	}

	const unsupportedOps = linux.MS_REMOUNT | linux.MS_BIND |
		linux.MS_SHARED | linux.MS_PRIVATE | linux.MS_SLAVE |
		linux.MS_UNBINDABLE | linux.MS_MOVE

	// Silently allow MS_NOSUID, since we don't implement set-id bits
	// anyway.
	const unsupportedFlags = linux.MS_NODEV | linux.MS_NOEXEC |
		linux.MS_NODIRATIME | linux.MS_STRICTATIME

	// Linux just allows passing any flags to mount(2) - it won't fail when
	// unknown or unsupported flags are passed. Since we don't implement
	// everything, we fail explicitly on flags that are unimplemented.
	if flags&(unsupportedOps|unsupportedFlags) != 0 {
		return 0, nil, syserror.EINVAL
	}

	rsys, ok := fs.FindFilesystem(fsType)
	if !ok {
		return 0, nil, syserror.ENODEV
	}
	if !rsys.AllowUserMount() {
		return 0, nil, syserror.EPERM
	}

	var superFlags fs.MountSourceFlags
	if flags&linux.MS_NOATIME == linux.MS_NOATIME {
		superFlags.NoAtime = true
	}
	if flags&linux.MS_RDONLY == linux.MS_RDONLY {
		superFlags.ReadOnly = true
	}

	rootInode, err := rsys.Mount(t, sourcePath, superFlags, data)
	if err != nil {
		return 0, nil, syserror.EINVAL
	}

	return 0, nil, fileOpOn(t, linux.AT_FDCWD, targetPath, true /* resolve */, func(root *fs.Dirent, d *fs.Dirent) error {
		return t.MountNamespace().Mount(t, d, rootInode)
	})
}

// Umount2 implements Linux syscall umount2(2).
func Umount2(t *kernel.Task, args arch.SyscallArguments) (uintptr, *kernel.SyscallControl, error) {
	addr := args[0].Pointer()
	flags := args[1].Int()

	const unsupported = linux.MNT_FORCE | linux.MNT_EXPIRE
	if flags&unsupported != 0 {
		return 0, nil, syserror.EINVAL
	}

	path, _, err := copyInPath(t, addr, false /* allowEmpty */)
	if err != nil {
		return 0, nil, err
	}

	// Must have CAP_SYS_ADMIN in the mount namespace's associated user
	// namespace.
	//
	// Currently, this is always the init task's user namespace.
	if !t.HasCapabilityIn(linux.CAP_SYS_ADMIN, t.MountNamespace().UserNamespace()) {
		return 0, nil, syserror.EPERM
	}

	resolve := flags&linux.UMOUNT_NOFOLLOW != linux.UMOUNT_NOFOLLOW
	detachOnly := flags&linux.MNT_DETACH == linux.MNT_DETACH

	return 0, nil, fileOpOn(t, linux.AT_FDCWD, path, resolve, func(root *fs.Dirent, d *fs.Dirent) error {
		return t.MountNamespace().Unmount(t, d, detachOnly)
	})
}
<commit_msg>Allow NULL data in mount(2)<commit_after>// Copyright 2018 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package linux

import (
	"gvisor.googlesource.com/gvisor/pkg/abi/linux"
	"gvisor.googlesource.com/gvisor/pkg/sentry/arch"
	"gvisor.googlesource.com/gvisor/pkg/sentry/fs"
	"gvisor.googlesource.com/gvisor/pkg/sentry/kernel"
	"gvisor.googlesource.com/gvisor/pkg/sentry/usermem"
	"gvisor.googlesource.com/gvisor/pkg/syserror"
)

// Mount implements Linux syscall mount(2).
func Mount(t *kernel.Task, args arch.SyscallArguments) (uintptr, *kernel.SyscallControl, error) {
	sourceAddr := args[0].Pointer()
	targetAddr := args[1].Pointer()
	typeAddr := args[2].Pointer()
	flags := args[3].Uint64()
	dataAddr := args[4].Pointer()

	fsType, err := t.CopyInString(typeAddr, usermem.PageSize)
	if err != nil {
		return 0, nil, err
	}

	sourcePath, _, err := copyInPath(t, sourceAddr, true /* allowEmpty */)
	if err != nil {
		return 0, nil, err
	}

	targetPath, _, err := copyInPath(t, targetAddr, false /* allowEmpty */)
	if err != nil {
		return 0, nil, err
	}

	data := ""
	if dataAddr != 0 {
		// In Linux, a full page is always copied in regardless of null
		// character placement, and the address is passed to each file system.
		// Most file systems always treat this data as a string, though, and so
		// do all of the ones we implement.
		data, err = t.CopyInString(dataAddr, usermem.PageSize)
		if err != nil {
			return 0, nil, err
		}
	}

	// Ignore magic value that was required before Linux 2.4.
	if flags&linux.MS_MGC_MSK == linux.MS_MGC_VAL {
		flags = flags &^ linux.MS_MGC_MSK
	}

	// Must have CAP_SYS_ADMIN in the mount namespace's associated user
	// namespace.
	if !t.HasCapabilityIn(linux.CAP_SYS_ADMIN, t.MountNamespace().UserNamespace()) {
		return 0, nil, syserror.EPERM
	}

	const unsupportedOps = linux.MS_REMOUNT | linux.MS_BIND |
		linux.MS_SHARED | linux.MS_PRIVATE | linux.MS_SLAVE |
		linux.MS_UNBINDABLE | linux.MS_MOVE

	// Silently allow MS_NOSUID, since we don't implement set-id bits
	// anyway.
	const unsupportedFlags = linux.MS_NODEV | linux.MS_NOEXEC |
		linux.MS_NODIRATIME | linux.MS_STRICTATIME

	// Linux just allows passing any flags to mount(2) - it won't fail when
	// unknown or unsupported flags are passed. Since we don't implement
	// everything, we fail explicitly on flags that are unimplemented.
	if flags&(unsupportedOps|unsupportedFlags) != 0 {
		return 0, nil, syserror.EINVAL
	}

	rsys, ok := fs.FindFilesystem(fsType)
	if !ok {
		return 0, nil, syserror.ENODEV
	}
	if !rsys.AllowUserMount() {
		return 0, nil, syserror.EPERM
	}

	var superFlags fs.MountSourceFlags
	if flags&linux.MS_NOATIME == linux.MS_NOATIME {
		superFlags.NoAtime = true
	}
	if flags&linux.MS_RDONLY == linux.MS_RDONLY {
		superFlags.ReadOnly = true
	}

	rootInode, err := rsys.Mount(t, sourcePath, superFlags, data)
	if err != nil {
		return 0, nil, syserror.EINVAL
	}

	return 0, nil, fileOpOn(t, linux.AT_FDCWD, targetPath, true /* resolve */, func(root *fs.Dirent, d *fs.Dirent) error {
		return t.MountNamespace().Mount(t, d, rootInode)
	})
}

// Umount2 implements Linux syscall umount2(2).
func Umount2(t *kernel.Task, args arch.SyscallArguments) (uintptr, *kernel.SyscallControl, error) {
	addr := args[0].Pointer()
	flags := args[1].Int()

	const unsupported = linux.MNT_FORCE | linux.MNT_EXPIRE
	if flags&unsupported != 0 {
		return 0, nil, syserror.EINVAL
	}

	path, _, err := copyInPath(t, addr, false /* allowEmpty */)
	if err != nil {
		return 0, nil, err
	}

	// Must have CAP_SYS_ADMIN in the mount namespace's associated user
	// namespace.
	//
	// Currently, this is always the init task's user namespace.
	if !t.HasCapabilityIn(linux.CAP_SYS_ADMIN, t.MountNamespace().UserNamespace()) {
		return 0, nil, syserror.EPERM
	}

	resolve := flags&linux.UMOUNT_NOFOLLOW != linux.UMOUNT_NOFOLLOW
	detachOnly := flags&linux.MNT_DETACH == linux.MNT_DETACH

	return 0, nil, fileOpOn(t, linux.AT_FDCWD, path, resolve, func(root *fs.Dirent, d *fs.Dirent) error {
		return t.MountNamespace().Unmount(t, d, detachOnly)
	})
}
<|endoftext|>


<commit_before>package mrsa

import (
	"crypto"
	"crypto/rsa"
	"errors"
	"math/big"
)

// This file implements the actual signing of messages.

// A Session is an object that contains enough PartialDecryptors to
// fully sign a message using PKCS1.
type Session struct {
	PublicKey
	decryptors []PartialDecryptor
}

// decrypt performs a full RSA decryption using multiple PartialDecryptors
func (session *Session) decrypt(c *big.Int) (* big.Int, error) {

	var err error
	ms := make([]*big.Int, len(session.decryptors))

	for i, d := range(session.decryptors) {
		ms[i], err = d.PartialDecrypt(c)
		if err != nil {
			return nil, err
		}
	}
	m := session.finalizeDecrypt(ms...)

	// we don't trust our partial decryptors very much, so we must
	// verify that the resulting decryption is valid.
	c2 := session.encrypt(m)
	if c2.Cmp(c) != 0 {
		return nil, errors.New("mrsa: decryption used a wrong RSA key")
	}

	return m, nil
}

// Most of this code is copied from http://golang.org/src/pkg/crypto/rsa/pkcs1v15.go

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

// SignPKCS1v15 calculates the signature of hashed using RSASSA-PKCS1-V1_5-SIGN from RSA PKCS#1 v1.5.
// Note that hashed must be the result of hashing the input message using the
// given hash function.
func (session *Session) SignPKCS1v15(hash crypto.Hash, hashed []byte) (s []byte, err error) {
	hashLen, prefix, err := pkcs1v15HashInfo(hash, len(hashed))
	if err != nil {
		return
	}

	tLen := len(prefix) + hashLen
	k := (session.N.BitLen() + 7) / 8
	if k < tLen+11 {
		return nil, rsa.ErrMessageTooLong
	}

	// EM = 0x00 || 0x01 || PS || 0x00 || T
	em := make([]byte, k)
	em[1] = 1
	for i := 2; i < k-tLen-1; i++ {
		em[i] = 0xff
	}
	copy(em[k-tLen:k-hashLen], prefix)
	copy(em[k-hashLen:k], hashed)

	m := new(big.Int).SetBytes(em)
	c, err := session.decrypt(m)
	if err != nil {
		return
	}

	copyWithLeftPad(em, c.Bytes())
	s = em
	return
}

func pkcs1v15HashInfo(hash crypto.Hash, inLen int) (hashLen int, prefix []byte, err error) {
	hashLen = hash.Size()
	if inLen != hashLen {
		return 0, nil, errors.New("mrsa: input must be hashed message")
	}
	prefix, ok := hashPrefixes[hash]
	if !ok {
		return 0, nil, errors.New("mrsa: unsupported hash function")
	}
	return
}

// copyWithLeftPad copies src to the end of dest, padding with zero bytes as
// needed.
func copyWithLeftPad(dest, src []byte) {
	numPaddingBytes := len(dest) - len(src)
	for i := 0; i < numPaddingBytes; i++ {
		dest[i] = 0
	}
	copy(dest[numPaddingBytes:], src)
}

var hashPrefixes = map[crypto.Hash][]byte{
	crypto.MD5:       {0x30, 0x20, 0x30, 0x0c, 0x06, 0x08, 0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 0x02, 0x05, 0x05, 0x00, 0x04, 0x10},
	crypto.SHA1:      {0x30, 0x21, 0x30, 0x09, 0x06, 0x05, 0x2b, 0x0e, 0x03, 0x02, 0x1a, 0x05, 0x00, 0x04, 0x14},
	crypto.SHA224:    {0x30, 0x2d, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x04, 0x05, 0x00, 0x04, 0x1c},
	crypto.SHA256:    {0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05, 0x00, 0x04, 0x20},
	crypto.SHA384:    {0x30, 0x41, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x02, 0x05, 0x00, 0x04, 0x30},
	crypto.SHA512:    {0x30, 0x51, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03, 0x05, 0x00, 0x04, 0x40},
	crypto.MD5SHA1:   {}, // A special TLS case which doesn't use an ASN1 prefix.
	crypto.RIPEMD160: {0x30, 0x20, 0x30, 0x08, 0x06, 0x06, 0x28, 0xcf, 0x06, 0x03, 0x00, 0x31, 0x04, 0x14},
}
<commit_msg>Export Decryptors<commit_after>package mrsa

import (
	"crypto"
	"crypto/rsa"
	"errors"
	"math/big"
)

// This file implements the actual signing of messages.

// A Session is an object that contains enough PartialDecryptors to
// fully sign a message using PKCS1.
type Session struct {
	PublicKey
	Decryptors []PartialDecryptor
}

// decrypt performs a full RSA decryption using multiple PartialDecryptors
func (session *Session) decrypt(c *big.Int) (* big.Int, error) {

	var err error
	ms := make([]*big.Int, len(session.Decryptors))

	for i, d := range(session.Decryptors) {
		ms[i], err = d.PartialDecrypt(c)
		if err != nil {
			return nil, err
		}
	}
	m := session.finalizeDecrypt(ms...)

	// we don't trust our partial decryptors very much, so we must
	// verify that the resulting decryption is valid.
	c2 := session.encrypt(m)
	if c2.Cmp(c) != 0 {
		return nil, errors.New("mrsa: decryption used a wrong RSA key")
	}

	return m, nil
}

// Most of this code is copied from http://golang.org/src/pkg/crypto/rsa/pkcs1v15.go

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

// SignPKCS1v15 calculates the signature of hashed using RSASSA-PKCS1-V1_5-SIGN from RSA PKCS#1 v1.5.
// Note that hashed must be the result of hashing the input message using the
// given hash function.
func (session *Session) SignPKCS1v15(hash crypto.Hash, hashed []byte) (s []byte, err error) {
	hashLen, prefix, err := pkcs1v15HashInfo(hash, len(hashed))
	if err != nil {
		return
	}

	tLen := len(prefix) + hashLen
	k := (session.N.BitLen() + 7) / 8
	if k < tLen+11 {
		return nil, rsa.ErrMessageTooLong
	}

	// EM = 0x00 || 0x01 || PS || 0x00 || T
	em := make([]byte, k)
	em[1] = 1
	for i := 2; i < k-tLen-1; i++ {
		em[i] = 0xff
	}
	copy(em[k-tLen:k-hashLen], prefix)
	copy(em[k-hashLen:k], hashed)

	m := new(big.Int).SetBytes(em)
	c, err := session.decrypt(m)
	if err != nil {
		return
	}

	copyWithLeftPad(em, c.Bytes())
	s = em
	return
}

func pkcs1v15HashInfo(hash crypto.Hash, inLen int) (hashLen int, prefix []byte, err error) {
	hashLen = hash.Size()
	if inLen != hashLen {
		return 0, nil, errors.New("mrsa: input must be hashed message")
	}
	prefix, ok := hashPrefixes[hash]
	if !ok {
		return 0, nil, errors.New("mrsa: unsupported hash function")
	}
	return
}

// copyWithLeftPad copies src to the end of dest, padding with zero bytes as
// needed.
func copyWithLeftPad(dest, src []byte) {
	numPaddingBytes := len(dest) - len(src)
	for i := 0; i < numPaddingBytes; i++ {
		dest[i] = 0
	}
	copy(dest[numPaddingBytes:], src)
}

var hashPrefixes = map[crypto.Hash][]byte{
	crypto.MD5:       {0x30, 0x20, 0x30, 0x0c, 0x06, 0x08, 0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 0x02, 0x05, 0x05, 0x00, 0x04, 0x10},
	crypto.SHA1:      {0x30, 0x21, 0x30, 0x09, 0x06, 0x05, 0x2b, 0x0e, 0x03, 0x02, 0x1a, 0x05, 0x00, 0x04, 0x14},
	crypto.SHA224:    {0x30, 0x2d, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x04, 0x05, 0x00, 0x04, 0x1c},
	crypto.SHA256:    {0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05, 0x00, 0x04, 0x20},
	crypto.SHA384:    {0x30, 0x41, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x02, 0x05, 0x00, 0x04, 0x30},
	crypto.SHA512:    {0x30, 0x51, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03, 0x05, 0x00, 0x04, 0x40},
	crypto.MD5SHA1:   {}, // A special TLS case which doesn't use an ASN1 prefix.
	crypto.RIPEMD160: {0x30, 0x20, 0x30, 0x08, 0x06, 0x06, 0x28, 0xcf, 0x06, 0x03, 0x00, 0x31, 0x04, 0x14},
}
<|endoftext|>


<commit_before>package smux

import (
	"io"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
)

const (
	defaultAcceptBacklog = 1024
)

const (
	errBrokenPipe      = "broken pipe"
	errInvalidProtocol = "invalid protocol version"
)

// Session defines a multiplexed connection for streams
type Session struct {
	conn      io.ReadWriteCloser
	writeLock sync.Mutex

	config       *Config
	nextStreamID uint32 // next stream identifier

	bucket     int32      // token bucket
	bucketCond *sync.Cond // used for waiting for tokens

	streams    map[uint32]*Stream // all streams in this session
	streamLock sync.Mutex         // locks streams

	die            chan struct{} // flag session has died
	dieLock        sync.Mutex
	chAccepts      chan *Stream
	chClosedStream chan uint32

	dataReady int32 // flag data has arrived
}

func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
	s := new(Session)
	s.die = make(chan struct{})
	s.conn = conn
	s.config = config
	s.streams = make(map[uint32]*Stream)
	s.chAccepts = make(chan *Stream, defaultAcceptBacklog)
	s.bucket = int32(config.MaxReceiveBuffer)
	s.bucketCond = sync.NewCond(&sync.Mutex{})
	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 2
	}
	go s.recvLoop()
	go s.keepalive()
	return s
}

// OpenStream is used to create a new stream
func (s *Session) OpenStream() (*Stream, error) {
	if s.IsClosed() {
		return nil, errors.New(errBrokenPipe)
	}

	sid := atomic.AddUint32(&s.nextStreamID, 2)
	stream := newStream(sid, s.config.MaxFrameSize, s)

	s.streamLock.Lock()
	s.streams[sid] = stream
	s.streamLock.Unlock()

	s.writeFrame(newFrame(cmdSYN, sid))
	return stream, nil
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	select {
	case stream := <-s.chAccepts:
		return stream, nil
	case <-s.die:
		return nil, errors.New(errBrokenPipe)
	}
}

// Close is used to close the session and all streams.
func (s *Session) Close() error {
	s.dieLock.Lock()
	defer s.dieLock.Unlock()

	select {
	case <-s.die:
		return errors.New(errBrokenPipe)
	default:
		s.streamLock.Lock()
		for k := range s.streams {
			s.streams[k].sessionClose()
		}
		s.streamLock.Unlock()
		s.conn.Close()
		close(s.die)
		s.bucketCond.Signal()
	}
	return nil
}

// IsClosed does a safe check to see if we have shutdown
func (s *Session) IsClosed() bool {
	select {
	case <-s.die:
		return true
	default:
		return false
	}
}

// NumStreams returns the number of currently open streams
func (s *Session) NumStreams() int {
	if s.IsClosed() {
		return 0
	}
	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	return len(s.streams)
}

// notify the session that a stream has closed
func (s *Session) streamClosed(sid uint32) {
	s.streamLock.Lock()
	if n := s.streams[sid].recycleTokens(); n > 0 { // return remaining tokens to the bucket
		if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
			s.bucketCond.Signal()
		}
	}
	delete(s.streams, sid)
	s.streamLock.Unlock()
}

// returnTokens is called by stream to return token after read
func (s *Session) returnTokens(n int) {
	if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
		s.bucketCond.Signal()
	}
}

// session read a frame from underlying connection
// it's data is pointed to the input buffer
func (s *Session) readFrame(buffer []byte) (f Frame, err error) {
	if _, err := io.ReadFull(s.conn, buffer[:headerSize]); err != nil {
		return f, errors.Wrap(err, "readFrame")
	}

	dec := rawHeader(buffer)
	if dec.Version() != version {
		return f, errors.New(errInvalidProtocol)
	}

	if length := dec.Length(); length > 0 {
		if _, err := io.ReadFull(s.conn, buffer[headerSize:headerSize+length]); err != nil {
			return f, errors.Wrap(err, "readFrame")
		}
		f.zeroCopyUnmarshal(buffer[:headerSize+length])
		return f, nil
	}
	f.zeroCopyUnmarshal(buffer[:headerSize])
	return f, nil
}

// recvLoop keeps on reading from underlying connection if tokens are available
func (s *Session) recvLoop() {
	buffer := make([]byte, (1<<16)+headerSize)
	for {
		s.bucketCond.L.Lock()
		for atomic.LoadInt32(&s.bucket) <= 0 && !s.IsClosed() {
			s.bucketCond.Wait()
		}
		s.bucketCond.L.Unlock()

		if s.IsClosed() {
			return
		}

		if f, err := s.readFrame(buffer); err == nil {
			atomic.StoreInt32(&s.dataReady, 1)

			switch f.cmd {
			case cmdNOP:
			case cmdSYN:
				s.streamLock.Lock()
				if _, ok := s.streams[f.sid]; !ok {
					stream := newStream(f.sid, s.config.MaxFrameSize, s)
					s.streams[f.sid] = stream
					go func() {
						select {
						case s.chAccepts <- stream:
						case <-s.die:
						}
					}()
				}
				s.streamLock.Unlock()
			case cmdRST:
				s.streamLock.Lock()
				if stream, ok := s.streams[f.sid]; ok {
					stream.markRST()
					stream.notifyReadEvent()
				}
				s.streamLock.Unlock()
			case cmdPSH:
				s.streamLock.Lock()
				if stream, ok := s.streams[f.sid]; ok {
					atomic.AddInt32(&s.bucket, -int32(len(f.data)))
					stream.pushBytes(f.data)
					stream.notifyReadEvent()
				}
				s.streamLock.Unlock()
			default:
				s.Close()
				return
			}
		} else {
			s.Close()
			return
		}
	}
}

func (s *Session) keepalive() {
	tickerPing := time.NewTicker(s.config.KeepAliveInterval)
	tickerTimeout := time.NewTicker(s.config.KeepAliveTimeout)
	defer tickerPing.Stop()
	defer tickerTimeout.Stop()
	for {
		select {
		case <-tickerPing.C:
			s.writeFrame(newFrame(cmdNOP, 0))
		case <-tickerTimeout.C:
			if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) {
				s.Close()
				return
			}
		case <-s.die:
			return
		}
	}
}

// writeFrame writes the frame to the underlying connection
// and returns the number of bytes written if successful
func (s *Session) writeFrame(f Frame) (n int, err error) {
	bts, _ := f.MarshalBinary()
	s.writeLock.Lock()
	n, err = s.conn.Write(bts)
	s.writeLock.Unlock()
	return n, err
}

// writeBinary writes the byte slice to the underlying connection
func (s *Session) writeBinary(bts []byte) (n int, err error) {
	s.writeLock.Lock()
	n, err = s.conn.Write(bts)
	s.writeLock.Unlock()
	return n, err
}
<commit_msg>remove unused var<commit_after>package smux

import (
	"io"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
)

const (
	defaultAcceptBacklog = 1024
)

const (
	errBrokenPipe      = "broken pipe"
	errInvalidProtocol = "invalid protocol version"
)

// Session defines a multiplexed connection for streams
type Session struct {
	conn      io.ReadWriteCloser
	writeLock sync.Mutex

	config       *Config
	nextStreamID uint32 // next stream identifier

	bucket     int32      // token bucket
	bucketCond *sync.Cond // used for waiting for tokens

	streams    map[uint32]*Stream // all streams in this session
	streamLock sync.Mutex         // locks streams

	die       chan struct{} // flag session has died
	dieLock   sync.Mutex
	chAccepts chan *Stream

	dataReady int32 // flag data has arrived
}

func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
	s := new(Session)
	s.die = make(chan struct{})
	s.conn = conn
	s.config = config
	s.streams = make(map[uint32]*Stream)
	s.chAccepts = make(chan *Stream, defaultAcceptBacklog)
	s.bucket = int32(config.MaxReceiveBuffer)
	s.bucketCond = sync.NewCond(&sync.Mutex{})
	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 2
	}
	go s.recvLoop()
	go s.keepalive()
	return s
}

// OpenStream is used to create a new stream
func (s *Session) OpenStream() (*Stream, error) {
	if s.IsClosed() {
		return nil, errors.New(errBrokenPipe)
	}

	sid := atomic.AddUint32(&s.nextStreamID, 2)
	stream := newStream(sid, s.config.MaxFrameSize, s)

	s.streamLock.Lock()
	s.streams[sid] = stream
	s.streamLock.Unlock()

	s.writeFrame(newFrame(cmdSYN, sid))
	return stream, nil
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	select {
	case stream := <-s.chAccepts:
		return stream, nil
	case <-s.die:
		return nil, errors.New(errBrokenPipe)
	}
}

// Close is used to close the session and all streams.
func (s *Session) Close() error {
	s.dieLock.Lock()
	defer s.dieLock.Unlock()

	select {
	case <-s.die:
		return errors.New(errBrokenPipe)
	default:
		s.streamLock.Lock()
		for k := range s.streams {
			s.streams[k].sessionClose()
		}
		s.streamLock.Unlock()
		s.conn.Close()
		close(s.die)
		s.bucketCond.Signal()
	}
	return nil
}

// IsClosed does a safe check to see if we have shutdown
func (s *Session) IsClosed() bool {
	select {
	case <-s.die:
		return true
	default:
		return false
	}
}

// NumStreams returns the number of currently open streams
func (s *Session) NumStreams() int {
	if s.IsClosed() {
		return 0
	}
	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	return len(s.streams)
}

// notify the session that a stream has closed
func (s *Session) streamClosed(sid uint32) {
	s.streamLock.Lock()
	if n := s.streams[sid].recycleTokens(); n > 0 { // return remaining tokens to the bucket
		if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
			s.bucketCond.Signal()
		}
	}
	delete(s.streams, sid)
	s.streamLock.Unlock()
}

// returnTokens is called by stream to return token after read
func (s *Session) returnTokens(n int) {
	if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
		s.bucketCond.Signal()
	}
}

// session read a frame from underlying connection
// it's data is pointed to the input buffer
func (s *Session) readFrame(buffer []byte) (f Frame, err error) {
	if _, err := io.ReadFull(s.conn, buffer[:headerSize]); err != nil {
		return f, errors.Wrap(err, "readFrame")
	}

	dec := rawHeader(buffer)
	if dec.Version() != version {
		return f, errors.New(errInvalidProtocol)
	}

	if length := dec.Length(); length > 0 {
		if _, err := io.ReadFull(s.conn, buffer[headerSize:headerSize+length]); err != nil {
			return f, errors.Wrap(err, "readFrame")
		}
		f.zeroCopyUnmarshal(buffer[:headerSize+length])
		return f, nil
	}
	f.zeroCopyUnmarshal(buffer[:headerSize])
	return f, nil
}

// recvLoop keeps on reading from underlying connection if tokens are available
func (s *Session) recvLoop() {
	buffer := make([]byte, (1<<16)+headerSize)
	for {
		s.bucketCond.L.Lock()
		for atomic.LoadInt32(&s.bucket) <= 0 && !s.IsClosed() {
			s.bucketCond.Wait()
		}
		s.bucketCond.L.Unlock()

		if s.IsClosed() {
			return
		}

		if f, err := s.readFrame(buffer); err == nil {
			atomic.StoreInt32(&s.dataReady, 1)

			switch f.cmd {
			case cmdNOP:
			case cmdSYN:
				s.streamLock.Lock()
				if _, ok := s.streams[f.sid]; !ok {
					stream := newStream(f.sid, s.config.MaxFrameSize, s)
					s.streams[f.sid] = stream
					go func() {
						select {
						case s.chAccepts <- stream:
						case <-s.die:
						}
					}()
				}
				s.streamLock.Unlock()
			case cmdRST:
				s.streamLock.Lock()
				if stream, ok := s.streams[f.sid]; ok {
					stream.markRST()
					stream.notifyReadEvent()
				}
				s.streamLock.Unlock()
			case cmdPSH:
				s.streamLock.Lock()
				if stream, ok := s.streams[f.sid]; ok {
					atomic.AddInt32(&s.bucket, -int32(len(f.data)))
					stream.pushBytes(f.data)
					stream.notifyReadEvent()
				}
				s.streamLock.Unlock()
			default:
				s.Close()
				return
			}
		} else {
			s.Close()
			return
		}
	}
}

func (s *Session) keepalive() {
	tickerPing := time.NewTicker(s.config.KeepAliveInterval)
	tickerTimeout := time.NewTicker(s.config.KeepAliveTimeout)
	defer tickerPing.Stop()
	defer tickerTimeout.Stop()
	for {
		select {
		case <-tickerPing.C:
			s.writeFrame(newFrame(cmdNOP, 0))
		case <-tickerTimeout.C:
			if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) {
				s.Close()
				return
			}
		case <-s.die:
			return
		}
	}
}

// writeFrame writes the frame to the underlying connection
// and returns the number of bytes written if successful
func (s *Session) writeFrame(f Frame) (n int, err error) {
	bts, _ := f.MarshalBinary()
	s.writeLock.Lock()
	n, err = s.conn.Write(bts)
	s.writeLock.Unlock()
	return n, err
}

// writeBinary writes the byte slice to the underlying connection
func (s *Session) writeBinary(bts []byte) (n int, err error) {
	s.writeLock.Lock()
	n, err = s.conn.Write(bts)
	s.writeLock.Unlock()
	return n, err
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/subgraph/fw-daemon/proc"
	"github.com/yawning/bulb"

	"bufio"
	"fmt"
	"log"
	"net"
	"strconv"
	"strings"
	"sync"
)

// ProxyListener is used to listen for
// Tor Control port connections and
// dispatches them to worker sessions
// to implement the filtering proxy pipeline.
type ProxyListener struct {
	cfg         *RoflcoptorConfig
	watch       bool
	wg          *sync.WaitGroup
	tcpListener *net.TCPListener
	errChan     chan error
}

// NewProxyListener creates a new ProxyListener given
// a configuration structure.
func NewProxyListener(cfg *RoflcoptorConfig, wg *sync.WaitGroup, watch bool) (*ProxyListener, error) {
	p := ProxyListener{
		cfg:   cfg,
		wg:    wg,
		watch: watch,
	}
	return &p, nil
}

// InitAllListeners initialize all tor control port proxy listeners
func (p *ProxyListener) InitAllListeners() {
	// XXX TODO add UNIX domain socket listener
	p.wg.Add(1)
	p.FilterTCPAcceptLoop()
}

// FilterTCPAcceptLoop and listens and accepts new
// connections and passes them into our filter proxy pipeline
func (p *ProxyListener) FilterTCPAcceptLoop() {
	defer p.wg.Done()

	tcpAddr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%s", p.cfg.ListenIP, p.cfg.ListenTCPPort))
	if err != nil {
		panic(fmt.Sprintf("Failed to resolve TCP configured filter port: %s\n", err))
	}
	p.tcpListener, err = net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		panic(fmt.Sprintf("Failed to listen on the filter port: %s\n", err))
	}
	defer p.tcpListener.Close()

	// Listen for incoming connections, and dispatch workers.
	// XXX TODO use a select statement to implement a stop/shutdown channel?
	for {
		conn, err := p.tcpListener.Accept()
		if err != nil {
			if e, ok := err.(net.Error); ok && !e.Temporary() {
				log.Printf("ERR/tor: Failed to Accept(): %v", err)
			}
		}

		// Create the appropriate session instance.
		s := NewProxySession(p.cfg, conn, p.watch)
		go s.sessionWorker()
	}
}

// ProxySession is used to manage is single Tor Control port filter proxy session
type ProxySession struct {
	cfg                *RoflcoptorConfig
	clientFilterPolicy *FilterConfig
	serverFilterPolicy *FilterConfig

	appConn          net.Conn
	appConnReader    *bufio.Reader
	appConnWriteLock sync.Mutex

	torConn   *bulb.Conn
	protoInfo *bulb.ProtocolInfo

	watch   bool
	errChan chan error
	sync.WaitGroup
}

// NewProxySession creates a ProxySession given a client's connection
// to our proxy listener and a watch bool.
func NewProxySession(cfg *RoflcoptorConfig, conn net.Conn, watch bool) *ProxySession {
	s := &ProxySession{
		cfg:           cfg,
		watch:         watch,
		appConn:       conn,
		appConnReader: bufio.NewReader(conn),
		errChan:       make(chan error, 2),
	}
	return s
}

// getFilterPolicy returns a *ServerClientFilterConfig
// (session policy) if one can be found, otherwise nil is returned.
// Note that the connection is closed upon fatal error,
// when we fail to lookup the proc info, we return nil and an error.
func (s *ProxySession) getFilterPolicy() (*ServerClientFilterConfig, error) {

	fields := strings.Split(s.appConn.RemoteAddr().String(), ":")
	dstPortStr := fields[1]
	dstIP := net.ParseIP(s.cfg.ListenIP)

	if dstIP == nil {
		s.appConn.Close()
		return nil, fmt.Errorf("net.ParseIP fail for: %s\n", s.cfg.ListenIP)
	}

	srcP, _ := strconv.ParseUint(dstPortStr, 10, 16)
	dstP, _ := strconv.ParseUint(s.cfg.ListenTCPPort, 10, 16)
	procInfo := proc.LookupTCPSocketProcess(uint16(srcP), dstIP, uint16(dstP))
	log.Printf("proc exec path: %s\n", procInfo.ExePath)
	if procInfo == nil {
		s.appConn.Close()
		return nil, fmt.Errorf("Could not find process information for: %d %s %d\n", srcP, dstIP, dstP)
	}

	filter := getFilterForPathAndUID(procInfo.ExePath, procInfo.Uid)
	if filter == nil {
		filter = getFilterForPath(procInfo.ExePath)
	} else {
		log.Printf("No filters found for: %s (%d)\n", procInfo.ExePath, procInfo.Uid)
		filter = nil
	}
	return filter, nil
}

func (s *ProxySession) initTorControl() error {
	var err error

	// Connect to the real control port.
	if s.torConn, err = bulb.Dial("unix", s.cfg.TorControlSocketPath); err != nil {
		return fmt.Errorf("ERR/tor: Failed to connect to tor control port: %v", err)
	}

	// Issue a PROTOCOLINFO, so we can send a realistic response.
	if s.protoInfo, err = s.torConn.ProtocolInfo(); err != nil {
		s.torConn.Close()
		return fmt.Errorf("ERR/tor: Failed to issue PROTOCOLINFO: %v", err)
	}

	// Authenticate with the real tor control port.
	// XXX: Pull password out of `b.s.cfg`.
	if err = s.torConn.Authenticate(""); err != nil {
		s.torConn.Close()
		return fmt.Errorf("ERR/tor: Failed to authenticate: %v", err)
	}
	return nil
}

func (s *ProxySession) sessionWorker() {
	defer s.appConn.Close()

	var policy *ServerClientFilterConfig
	var err error

	clientAddr := s.appConn.RemoteAddr()
	log.Printf("INFO/tor: New ctrl connection from: %s", clientAddr)

	policy, err = s.getFilterPolicy()
	if err != nil {
		log.Printf("proc info query failure; connection from %s aborted: %s\n", clientAddr, err)
		return
	}
	if policy == nil {
		s.clientFilterPolicy = nil
		s.serverFilterPolicy = nil
	} else {
		s.clientFilterPolicy = &FilterConfig{
			Allowed:             policy.ClientAllowed,
			AllowedPrefixes:     policy.ClientAllowedPrefixes,
			Replacements:        policy.ClientReplacements,
			ReplacementPrefixes: policy.ClientReplacementPrefixes,
		}
		s.serverFilterPolicy = &FilterConfig{
			Allowed:             policy.ServerAllowed,
			AllowedPrefixes:     policy.ServerAllowedPrefixes,
			Replacements:        policy.ServerReplacements,
			ReplacementPrefixes: policy.ServerReplacementPrefixes,
		}
	}

	// Authenticate with the real control port
	err = s.initTorControl()
	defer s.torConn.Close()

	if err != nil {
		log.Printf("RoflcopTor: Failed to authenticate with the tor control port: %s\n", err)
		return
	}

	s.appConnReader = bufio.NewReader(s.appConn)
	s.Add(2)
	go s.proxyFilterTorToApp()
	go s.proxyFilterAppToTor()

	// Wait till all sessions are finished, log and return.
	s.Wait()
	s.torConn.Close()
	s.appConn.Close()

	if len(s.errChan) > 0 {
		err := <-s.errChan
		log.Printf("INFO/tor: Closed client connection from: %s: %v", clientAddr, err)
	} else {
		log.Printf("INFO/tor: Closed client connection from: %v", clientAddr)
	}
}

func (s *ProxySession) proxyFilterTorToApp() {
	defer s.Done()

	var appConnLock sync.Mutex
	writeAppConn := func(b []byte) (int, error) {
		appConnLock.Lock() // XXX is the lock really needed here?
		defer appConnLock.Unlock()
		return s.appConn.Write(b)
	}

	rd := bufio.NewReader(s.torConn)
	for {
		line, err := rd.ReadBytes('\n')
		if err != nil {
			s.errChan <- err
			break
		}
		lineStr := strings.TrimSpace(string(line))
		log.Printf("A<-T: [%s]\n", lineStr)
		if s.watch {
			if _, err = writeAppConn([]byte(lineStr + "\n")); err != nil {
				s.errChan <- err
				break
			}
		} else {
			filterCommand(lineStr, "250 OK", writeAppConn, s.errChan, s.serverFilterPolicy)
		}
	}
}

func (s *ProxySession) proxyFilterAppToTor() {
	defer s.Done()

	writeToTor := func(line []byte) (int, error) {
		n, err := s.torConn.Write([]byte(line))
		return n, err
	}

	for {
		line, err := s.appConnReader.ReadBytes('\n')
		if err != nil {
			s.errChan <- err
			break
		}
		lineStr := strings.TrimSpace(string(line))
		log.Printf("A->T: [%s]\n", lineStr)

		if s.watch {
			_, err = writeToTor([]byte(lineStr + "\n"))
			if err != nil {
				s.errChan <- err
			}
		} else {
			filterCommand(lineStr, "", writeToTor, s.errChan, s.clientFilterPolicy)
		}
	}
}
<commit_msg>Migrated to extracted procsnitch library.<commit_after>package main

import (
	"bufio"
	"fmt"
	"log"
	"net"
	"strconv"
	"strings"
	"sync"

	"github.com/subgraph/go-procsnitch"
	"github.com/yawning/bulb"
)

// ProxyListener is used to listen for
// Tor Control port connections and
// dispatches them to worker sessions
// to implement the filtering proxy pipeline.
type ProxyListener struct {
	cfg         *RoflcoptorConfig
	watch       bool
	wg          *sync.WaitGroup
	tcpListener *net.TCPListener
	errChan     chan error
}

// NewProxyListener creates a new ProxyListener given
// a configuration structure.
func NewProxyListener(cfg *RoflcoptorConfig, wg *sync.WaitGroup, watch bool) (*ProxyListener, error) {
	p := ProxyListener{
		cfg:   cfg,
		wg:    wg,
		watch: watch,
	}
	return &p, nil
}

// InitAllListeners initialize all tor control port proxy listeners
func (p *ProxyListener) InitAllListeners() {
	// XXX TODO add UNIX domain socket listener
	p.wg.Add(1)
	p.FilterTCPAcceptLoop()
}

// FilterTCPAcceptLoop and listens and accepts new
// connections and passes them into our filter proxy pipeline
func (p *ProxyListener) FilterTCPAcceptLoop() {
	defer p.wg.Done()

	tcpAddr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%s", p.cfg.ListenIP, p.cfg.ListenTCPPort))
	if err != nil {
		panic(fmt.Sprintf("Failed to resolve TCP configured filter port: %s\n", err))
	}
	p.tcpListener, err = net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		panic(fmt.Sprintf("Failed to listen on the filter port: %s\n", err))
	}
	defer p.tcpListener.Close()

	// Listen for incoming connections, and dispatch workers.
	// XXX TODO use a select statement to implement a stop/shutdown channel?
	for {
		conn, err := p.tcpListener.Accept()
		if err != nil {
			if e, ok := err.(net.Error); ok && !e.Temporary() {
				log.Printf("ERR/tor: Failed to Accept(): %v", err)
			}
		}

		// Create the appropriate session instance.
		s := NewProxySession(p.cfg, conn, p.watch)
		go s.sessionWorker()
	}
}

// ProxySession is used to manage is single Tor Control port filter proxy session
type ProxySession struct {
	cfg                *RoflcoptorConfig
	clientFilterPolicy *FilterConfig
	serverFilterPolicy *FilterConfig

	appConn          net.Conn
	appConnReader    *bufio.Reader
	appConnWriteLock sync.Mutex

	torConn   *bulb.Conn
	protoInfo *bulb.ProtocolInfo

	watch   bool
	errChan chan error
	sync.WaitGroup
}

// NewProxySession creates a ProxySession given a client's connection
// to our proxy listener and a watch bool.
func NewProxySession(cfg *RoflcoptorConfig, conn net.Conn, watch bool) *ProxySession {
	s := &ProxySession{
		cfg:           cfg,
		watch:         watch,
		appConn:       conn,
		appConnReader: bufio.NewReader(conn),
		errChan:       make(chan error, 2),
	}
	return s
}

// getFilterPolicy returns a *ServerClientFilterConfig
// (session policy) if one can be found, otherwise nil is returned.
// Note that the connection is closed upon fatal error,
// when we fail to lookup the proc info, we return nil and an error.
func (s *ProxySession) getFilterPolicy() (*ServerClientFilterConfig, error) {

	fields := strings.Split(s.appConn.RemoteAddr().String(), ":")
	dstPortStr := fields[1]
	dstIP := net.ParseIP(s.cfg.ListenIP)

	if dstIP == nil {
		s.appConn.Close()
		return nil, fmt.Errorf("net.ParseIP fail for: %s\n", s.cfg.ListenIP)
	}

	srcP, _ := strconv.ParseUint(dstPortStr, 10, 16)
	dstP, _ := strconv.ParseUint(s.cfg.ListenTCPPort, 10, 16)
	procInfo := procsnitch.LookupTCPSocketProcess(uint16(srcP), dstIP, uint16(dstP))
	log.Printf("proc exec path: %s\n", procInfo.ExePath)
	if procInfo == nil {
		s.appConn.Close()
		return nil, fmt.Errorf("Could not find process information for: %d %s %d\n", srcP, dstIP, dstP)
	}

	filter := getFilterForPathAndUID(procInfo.ExePath, procInfo.UID)
	if filter == nil {
		filter = getFilterForPath(procInfo.ExePath)
	} else {
		log.Printf("No filters found for: %s (%d)\n", procInfo.ExePath, procInfo.UID)
		filter = nil
	}
	return filter, nil
}

func (s *ProxySession) initTorControl() error {
	var err error

	// Connect to the real control port.
	if s.torConn, err = bulb.Dial("unix", s.cfg.TorControlSocketPath); err != nil {
		return fmt.Errorf("ERR/tor: Failed to connect to tor control port: %v", err)
	}

	// Issue a PROTOCOLINFO, so we can send a realistic response.
	if s.protoInfo, err = s.torConn.ProtocolInfo(); err != nil {
		s.torConn.Close()
		return fmt.Errorf("ERR/tor: Failed to issue PROTOCOLINFO: %v", err)
	}

	// Authenticate with the real tor control port.
	// XXX: Pull password out of `b.s.cfg`.
	if err = s.torConn.Authenticate(""); err != nil {
		s.torConn.Close()
		return fmt.Errorf("ERR/tor: Failed to authenticate: %v", err)
	}
	return nil
}

func (s *ProxySession) sessionWorker() {
	defer s.appConn.Close()

	var policy *ServerClientFilterConfig
	var err error

	clientAddr := s.appConn.RemoteAddr()
	log.Printf("INFO/tor: New ctrl connection from: %s", clientAddr)

	policy, err = s.getFilterPolicy()
	if err != nil {
		log.Printf("proc info query failure; connection from %s aborted: %s\n", clientAddr, err)
		return
	}
	if policy == nil {
		s.clientFilterPolicy = nil
		s.serverFilterPolicy = nil
	} else {
		s.clientFilterPolicy = &FilterConfig{
			Allowed:             policy.ClientAllowed,
			AllowedPrefixes:     policy.ClientAllowedPrefixes,
			Replacements:        policy.ClientReplacements,
			ReplacementPrefixes: policy.ClientReplacementPrefixes,
		}
		s.serverFilterPolicy = &FilterConfig{
			Allowed:             policy.ServerAllowed,
			AllowedPrefixes:     policy.ServerAllowedPrefixes,
			Replacements:        policy.ServerReplacements,
			ReplacementPrefixes: policy.ServerReplacementPrefixes,
		}
	}

	// Authenticate with the real control port
	err = s.initTorControl()
	defer s.torConn.Close()

	if err != nil {
		log.Printf("RoflcopTor: Failed to authenticate with the tor control port: %s\n", err)
		return
	}

	s.appConnReader = bufio.NewReader(s.appConn)
	s.Add(2)
	go s.proxyFilterTorToApp()
	go s.proxyFilterAppToTor()

	// Wait till all sessions are finished, log and return.
	s.Wait()
	s.torConn.Close()
	s.appConn.Close()

	if len(s.errChan) > 0 {
		err := <-s.errChan
		log.Printf("INFO/tor: Closed client connection from: %s: %v", clientAddr, err)
	} else {
		log.Printf("INFO/tor: Closed client connection from: %v", clientAddr)
	}
}

func (s *ProxySession) proxyFilterTorToApp() {
	defer s.Done()

	var appConnLock sync.Mutex
	writeAppConn := func(b []byte) (int, error) {
		appConnLock.Lock() // XXX is the lock really needed here?
		defer appConnLock.Unlock()
		return s.appConn.Write(b)
	}

	rd := bufio.NewReader(s.torConn)
	for {
		line, err := rd.ReadBytes('\n')
		if err != nil {
			s.errChan <- err
			break
		}
		lineStr := strings.TrimSpace(string(line))
		log.Printf("A<-T: [%s]\n", lineStr)
		if s.watch {
			if _, err = writeAppConn([]byte(lineStr + "\n")); err != nil {
				s.errChan <- err
				break
			}
		} else {
			filterCommand(lineStr, "250 OK", writeAppConn, s.errChan, s.serverFilterPolicy)
		}
	}
}

func (s *ProxySession) proxyFilterAppToTor() {
	defer s.Done()

	writeToTor := func(line []byte) (int, error) {
		n, err := s.torConn.Write([]byte(line))
		return n, err
	}

	for {
		line, err := s.appConnReader.ReadBytes('\n')
		if err != nil {
			s.errChan <- err
			break
		}
		lineStr := strings.TrimSpace(string(line))
		log.Printf("A->T: [%s]\n", lineStr)

		if s.watch {
			_, err = writeToTor([]byte(lineStr + "\n"))
			if err != nil {
				s.errChan <- err
			}
		} else {
			filterCommand(lineStr, "", writeToTor, s.errChan, s.clientFilterPolicy)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      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 database

import (
	"time"
)

type Config struct {
	Name               string        `envconfig:"DB_NAME"`
	User               string        `envconfig:"DB_USER"`
	Host               string        `envconfig:"DB_HOST" default:"localhost"`
	Port               string        `envconfig:"DB_PORT" default:"5432"`
	SSLMode            string        `envconfig:"DB_SSLMODE" default:"required"`
	ConnectionTimeout  int           `envconfig:"DB_CONNECT_TIMEOUT"`
	Password           string        `envconfig:"DB_PASSWORD"`
	SSLCertPath        string        `envconfig:"DB_SSLCERT"`
	SSLKeyPath         string        `envconfig:"DB_SSLKEY"`
	SSLRootCertPath    string        `envconfig:"DB_SSLROOTCERT"`
	PoolMinConnections string        `envconfig:"DB_POOL_MIN_CONNS"`
	PoolMaxConnections string        `envconfig:"DB_POOL_MAX_CONNS"`
	PoolMaxConnLife    time.Duration `envconfig:"DB_POOL_MAX_CONN_LIFETIME"`
	PoolMaxConnIdle    time.Duration `envconfig:"DB_POOL_MAX_CONN_IDLE_TIME"`
	PoolHealthCheck    time.Duration `envconfig:"DB_POOL_HEALTH_CHECK_PERIOD"`
}

func (c *Config) DB() *Config {
	return c
}
<commit_msg>Use correct default SSL mode (#448)<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      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 database

import (
	"time"
)

type Config struct {
	Name               string        `envconfig:"DB_NAME"`
	User               string        `envconfig:"DB_USER"`
	Host               string        `envconfig:"DB_HOST" default:"localhost"`
	Port               string        `envconfig:"DB_PORT" default:"5432"`
	SSLMode            string        `envconfig:"DB_SSLMODE" default:"require"`
	ConnectionTimeout  int           `envconfig:"DB_CONNECT_TIMEOUT"`
	Password           string        `envconfig:"DB_PASSWORD"`
	SSLCertPath        string        `envconfig:"DB_SSLCERT"`
	SSLKeyPath         string        `envconfig:"DB_SSLKEY"`
	SSLRootCertPath    string        `envconfig:"DB_SSLROOTCERT"`
	PoolMinConnections string        `envconfig:"DB_POOL_MIN_CONNS"`
	PoolMaxConnections string        `envconfig:"DB_POOL_MAX_CONNS"`
	PoolMaxConnLife    time.Duration `envconfig:"DB_POOL_MAX_CONN_LIFETIME"`
	PoolMaxConnIdle    time.Duration `envconfig:"DB_POOL_MAX_CONN_IDLE_TIME"`
	PoolHealthCheck    time.Duration `envconfig:"DB_POOL_HEALTH_CHECK_PERIOD"`
}

func (c *Config) DB() *Config {
	return c
}
<|endoftext|>