<commit_before>package drouter

import (
	"fmt"
	log "github.com/Sirupsen/logrus"
	dockertypes "github.com/docker/engine-api/types"
	"github.com/vishvananda/netlink"
	"golang.org/x/net/context"
	"net"
	"strings"
)

type container struct {
	handle *netlink.Handle
}

func newContainerFromID(id string) (*container, error) {
	cjson, err := dockerClient.ContainerInspect(context.Background(), id)
	if err != nil {
		log.Error("Failed to inspect container with id: %v", id)
		return nil, err
	}
	log.Debugf("Inspected container: %v", cjson.Name)

	if !cjson.State.Running {
		return &container{}, nil
	}

	ch, err := netlinkHandleFromPid(cjson.State.Pid)
	if err != nil {
		return nil, err
	}

	return &container{handle: ch}, nil
}

//adds all known routes for provided container
func (c *container) addAllRoutes() error {
	//Loop through all static routes, ensure each one is installed in the container
	log.Info("Syncing static routes.")
	//add routes for all the static routes
	for _, sr := range staticRoutes {
		go c.addRoute(sr)
	}

	//Loop through all discovered networks, ensure each one is installed in the container
	//Unless it is covered by a static route already
	log.Info("Syncing discovered routes.")
	routes, err := netlink.RouteList(nil, netlink.FAMILY_ALL)
	if err != nil {
		log.Error("Failed to get my routes.")
		log.Error(err)
	}

Subnets:
	for _, r := range routes {
		if r.Gw != nil {
			continue
		}
		if localShortcut {
			if subnetEqualSubnet(r.Dst, networkID(p2p.network)) {
				continue
			}
		}
		for _, sr := range staticRoutes {
			if subnetContainsSubnet(sr, r.Dst) {
				log.Debugf("Skipping route %v covered by %v.", r.Dst, sr)
				break Subnets
			}
		}
		//TODO: filter routes that container is already connected to
		go c.addRoute(r.Dst)
	}

	return nil
}

func (c *container) addRoute(sn *net.IPNet) {
	gateway, err := c.getPathIP()
	if err != nil {
		log.Error(err)
		return
	}

	if (sn.IP.To4() == nil) != (gateway.To4() == nil) {
		// Dst is a different IP family
		return
	}

	route := &netlink.Route{
		Dst: sn,
		Gw:  gateway,
	}

	log.Infof("Adding route to %v via %v.", sn, gateway)
	err = c.handle.RouteAdd(route)
	if err != nil {
		if !strings.Contains(err.Error(), "file exists") {
			log.Error(err)
			return
		}
		routes, err2 := c.handle.RouteGet(sn.IP)
		if err2 != nil {
			log.Errorf("Failed to get container routes to: %v.", sn.IP)
			log.Error(err)
			return
		}
		for _, r := range routes {
			if r.Gw.Equal(gateway) || r.Gw == nil {
				return
			}
			err := c.handle.RouteDel(&netlink.Route{Dst: sn, Gw: r.Gw})
			if err != nil {
				log.Errorf("Failed to delete route to %v via %v.", sn, r.Gw)
				log.Error(err)
				return
			}
		}
		c.addRoute(sn)
	}
}

func (c *container) delRoutes(to, via *net.IPNet) {
	//get all container routes
	routes, err := c.getRoutes()
	if err != nil {
		log.Error("Failed to get container route table.")
		log.Error(err)
		return
	}

	//get all drouter ips
	ips, err := netlink.AddrList(nil, netlink.FAMILY_ALL)
	if err != nil {
		log.Error("Failed to get drouter ip addresses.")
		log.Error(err)
		return
	}

	for _, r := range routes {
		if r.Dst == nil {
			continue
		}
		if to != nil && !subnetContainsSubnet(to, r.Dst) {
			continue
		}

		for _, ipaddr := range ips {
			if !r.Gw.Equal(ipaddr.IP) {
				continue
			}
			if via != nil && !via.Contains(r.Gw) {
				continue
			}
			err := c.handle.RouteDel(&r)
			if err != nil {
				log.Errorf("Failed to delete container route to %v via %v", r.Dst, r.Gw)
				continue
			}
		}
	}
}

//sets the provided container's default route to the provided gateway
func (c *container) replaceGateway(gateway net.IP) error {
	log.Debugf("container.replaceGateway(%v)", gateway)

	var defr *netlink.Route
	//replace containers default gateway with drouter
	routes, err := c.getRoutes()
	if err != nil {
		return err
	}
	log.Debugf("container routes: %v", routes)
	for _, r := range routes {
		if r.Dst != nil {
			// Not the container gateway
			continue
		}

		defr = &r
	}

	//bail if the container gateway is already set to gateway
	if gateway.Equal(defr.Gw) {
		return nil
	}

	log.Debugf("Remove existing default route: %v", defr)
	err = c.handle.RouteDel(defr)
	if err != nil {
		return err
	}

	if gateway == nil || gateway.Equal(net.IP{}) {
		return nil
	}

	defr.Gw = gateway
	err = c.handle.RouteAdd(defr)
	if err != nil {
		return err
	}
	log.Debugf("Default route changed to: %v", defr)

	return nil
}

// called during a network connect event
func (c *container) connectEvent(drn *network) error {
	if !drn.isConnected() {
		connectWG.Add(1)
		go func() {
			defer connectWG.Done()
			drn.connect()
		}()
		//return now because the routeEvent will trigger routes to be installed in this container
		return nil
	}

	//let's push our routes into this new container
	gateway, err := c.getPathIP()
	if err != nil {
		log.Error("Failed to get container path IP.")
		return err
	}

	if localGateway {
		go c.replaceGateway(gateway)
	} else {
		go c.addAllRoutes()
	}
	return nil
}

// called when we detect a container has disconnected from a drouter network
func (c *container) disconnectEvent(drn *network) error {
	for _, ic := range drn.IPAM.Config {
		subnet, err := netlink.ParseIPNet(ic.Subnet)
		if err != nil {
			log.Errorf("Failed to parse ipam config subnet: %v", ic.Subnet)
			log.Error(err)
			continue
		}
		c.delRoutes(nil, subnet)
	}

	if _, err := c.getPathIP(); err == nil {
		c.addAllRoutes()
	}

	if aggressive {
		return nil
	}

	//if not aggressive mode, then we disconnect from the network if this is the last connected container

	//loop through all the containers
	dockerContainers, err := dockerClient.ContainerList(context.Background(), dockertypes.ContainerListOptions{})
	if err != nil {
		return err
	}
	for _, dc := range dockerContainers {
		if dc.HostConfig.NetworkMode == "host" {
			continue
		}
		if dc.ID == selfContainerID {
			continue
		}

		for n, _ := range dc.NetworkSettings.Networks {
			if n == drn.Name {
				// This netowork is still in use
				return nil
			}
		}
	}

	select {
	case _ = <-stopChan:
	default:
		go drn.disconnect()
	}

	return nil
}

//returns a drouter IP that is on some same network as provided container
func (c *container) getPathIP() (net.IP, error) {
	if c.handle == nil {
		return nil, fmt.Errorf("No namespace handle for container, is it running?")
	}

	addrs, err := c.handle.AddrList(nil, netlink.FAMILY_V4)
	if err != nil {
		log.Error("Failed to list container addresses.")
		return nil, err
	}

	for _, addr := range addrs {
		if addr.Label == "lo" {
			continue
		}

		log.Debugf("Getting my route to container IP: %v", addr.IP)
		srcRoutes, err := netlink.RouteGet(addr.IP)
		if err != nil {
			log.Error(err)
			continue
		}
		for _, srcRoute := range srcRoutes {
			if srcRoute.Gw == nil {
				return srcRoute.Src, nil
			}
		}
	}

	return nil, fmt.Errorf("No direct connection to container.")
}

func (c *container) getRoutes() ([]netlink.Route, error) {
	if c.handle != nil {
		return c.handle.RouteList(nil, netlink.FAMILY_ALL)
	}

	return make([]netlink.Route, 0), nil
}
<commit_msg>wait for connections before disconnect last container<commit_after>package drouter

import (
	"fmt"
	log "github.com/Sirupsen/logrus"
	dockertypes "github.com/docker/engine-api/types"
	"github.com/vishvananda/netlink"
	"golang.org/x/net/context"
	"net"
	"strings"
)

type container struct {
	handle *netlink.Handle
}

func newContainerFromID(id string) (*container, error) {
	cjson, err := dockerClient.ContainerInspect(context.Background(), id)
	if err != nil {
		log.Error("Failed to inspect container with id: %v", id)
		return nil, err
	}
	log.Debugf("Inspected container: %v", cjson.Name)

	if !cjson.State.Running {
		return &container{}, nil
	}

	ch, err := netlinkHandleFromPid(cjson.State.Pid)
	if err != nil {
		return nil, err
	}

	return &container{handle: ch}, nil
}

//adds all known routes for provided container
func (c *container) addAllRoutes() error {
	//Loop through all static routes, ensure each one is installed in the container
	log.Info("Syncing static routes.")
	//add routes for all the static routes
	for _, sr := range staticRoutes {
		go c.addRoute(sr)
	}

	//Loop through all discovered networks, ensure each one is installed in the container
	//Unless it is covered by a static route already
	log.Info("Syncing discovered routes.")
	routes, err := netlink.RouteList(nil, netlink.FAMILY_ALL)
	if err != nil {
		log.Error("Failed to get my routes.")
		log.Error(err)
	}

Subnets:
	for _, r := range routes {
		if r.Gw != nil {
			continue
		}
		if localShortcut {
			if subnetEqualSubnet(r.Dst, networkID(p2p.network)) {
				continue
			}
		}
		for _, sr := range staticRoutes {
			if subnetContainsSubnet(sr, r.Dst) {
				log.Debugf("Skipping route %v covered by %v.", r.Dst, sr)
				break Subnets
			}
		}
		//TODO: filter routes that container is already connected to
		go c.addRoute(r.Dst)
	}

	return nil
}

func (c *container) addRoute(sn *net.IPNet) {
	gateway, err := c.getPathIP()
	if err != nil {
		log.Error(err)
		return
	}

	if (sn.IP.To4() == nil) != (gateway.To4() == nil) {
		// Dst is a different IP family
		return
	}

	route := &netlink.Route{
		Dst: sn,
		Gw:  gateway,
	}

	log.Infof("Adding route to %v via %v.", sn, gateway)
	err = c.handle.RouteAdd(route)
	if err != nil {
		if !strings.Contains(err.Error(), "file exists") {
			log.Error(err)
			return
		}
		routes, err2 := c.handle.RouteGet(sn.IP)
		if err2 != nil {
			log.Errorf("Failed to get container routes to: %v.", sn.IP)
			log.Error(err)
			return
		}
		for _, r := range routes {
			if r.Gw.Equal(gateway) || r.Gw == nil {
				return
			}
			err := c.handle.RouteDel(&netlink.Route{Dst: sn, Gw: r.Gw})
			if err != nil {
				log.Errorf("Failed to delete route to %v via %v.", sn, r.Gw)
				log.Error(err)
				return
			}
		}
		c.addRoute(sn)
	}
}

func (c *container) delRoutes(to, via *net.IPNet) {
	//get all container routes
	routes, err := c.getRoutes()
	if err != nil {
		log.Error("Failed to get container route table.")
		log.Error(err)
		return
	}

	//get all drouter ips
	ips, err := netlink.AddrList(nil, netlink.FAMILY_ALL)
	if err != nil {
		log.Error("Failed to get drouter ip addresses.")
		log.Error(err)
		return
	}

	for _, r := range routes {
		if r.Dst == nil {
			continue
		}
		if to != nil && !subnetContainsSubnet(to, r.Dst) {
			continue
		}

		for _, ipaddr := range ips {
			if !r.Gw.Equal(ipaddr.IP) {
				continue
			}
			if via != nil && !via.Contains(r.Gw) {
				continue
			}
			err := c.handle.RouteDel(&r)
			if err != nil {
				log.Errorf("Failed to delete container route to %v via %v", r.Dst, r.Gw)
				continue
			}
		}
	}
}

//sets the provided container's default route to the provided gateway
func (c *container) replaceGateway(gateway net.IP) error {
	log.Debugf("container.replaceGateway(%v)", gateway)

	var defr *netlink.Route
	//replace containers default gateway with drouter
	routes, err := c.getRoutes()
	if err != nil {
		return err
	}
	log.Debugf("container routes: %v", routes)
	for _, r := range routes {
		if r.Dst != nil {
			// Not the container gateway
			continue
		}

		defr = &r
	}

	//bail if the container gateway is already set to gateway
	if gateway.Equal(defr.Gw) {
		return nil
	}

	log.Debugf("Remove existing default route: %v", defr)
	err = c.handle.RouteDel(defr)
	if err != nil {
		return err
	}

	if gateway == nil || gateway.Equal(net.IP{}) {
		return nil
	}

	defr.Gw = gateway
	err = c.handle.RouteAdd(defr)
	if err != nil {
		return err
	}
	log.Debugf("Default route changed to: %v", defr)

	return nil
}

// called during a network connect event
func (c *container) connectEvent(drn *network) error {
	if !drn.isConnected() {
		connectWG.Add(1)
		go func() {
			defer connectWG.Done()
			drn.connect()
		}()
		//return now because the routeEvent will trigger routes to be installed in this container
		return nil
	}

	//let's push our routes into this new container
	gateway, err := c.getPathIP()
	if err != nil {
		log.Error("Failed to get container path IP.")
		return err
	}

	if localGateway {
		go c.replaceGateway(gateway)
	} else {
		go c.addAllRoutes()
	}
	return nil
}

// called when we detect a container has disconnected from a drouter network
func (c *container) disconnectEvent(drn *network) error {
	for _, ic := range drn.IPAM.Config {
		subnet, err := netlink.ParseIPNet(ic.Subnet)
		if err != nil {
			log.Errorf("Failed to parse ipam config subnet: %v", ic.Subnet)
			log.Error(err)
			continue
		}
		c.delRoutes(nil, subnet)
	}

	if _, err := c.getPathIP(); err == nil {
		c.addAllRoutes()
	}

	if aggressive {
		return nil
	}

	//if not aggressive mode, then we disconnect from the network if this is the last connected container
	connectWG.Wait()

	//loop through all the containers
	dockerContainers, err := dockerClient.ContainerList(context.Background(), dockertypes.ContainerListOptions{})
	if err != nil {
		return err
	}
	for _, dc := range dockerContainers {
		if dc.HostConfig.NetworkMode == "host" {
			continue
		}
		if dc.ID == selfContainerID {
			continue
		}

		for n, _ := range dc.NetworkSettings.Networks {
			if n == drn.Name {
				// This netowork is still in use
				return nil
			}
		}
	}

	select {
	case _ = <-stopChan:
	default:
		go drn.disconnect()
	}

	return nil
}

//returns a drouter IP that is on some same network as provided container
func (c *container) getPathIP() (net.IP, error) {
	if c.handle == nil {
		return nil, fmt.Errorf("No namespace handle for container, is it running?")
	}

	addrs, err := c.handle.AddrList(nil, netlink.FAMILY_V4)
	if err != nil {
		log.Error("Failed to list container addresses.")
		return nil, err
	}

	for _, addr := range addrs {
		if addr.Label == "lo" {
			continue
		}

		log.Debugf("Getting my route to container IP: %v", addr.IP)
		srcRoutes, err := netlink.RouteGet(addr.IP)
		if err != nil {
			log.Error(err)
			continue
		}
		for _, srcRoute := range srcRoutes {
			if srcRoute.Gw == nil {
				return srcRoute.Src, nil
			}
		}
	}

	return nil, fmt.Errorf("No direct connection to container.")
}

func (c *container) getRoutes() ([]netlink.Route, error) {
	if c.handle != nil {
		return c.handle.RouteList(nil, netlink.FAMILY_ALL)
	}

	return make([]netlink.Route, 0), nil
}
<|endoftext|>


<commit_before>package vizzini_test

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"
	"time"

	"code.cloudfoundry.org/bbs/models"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/routing-info/cfroutes"

	. "code.cloudfoundry.org/vizzini/matchers"
	"github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

const HealthyCheckInterval = 30 * time.Second
const ConvergerInterval = 30 * time.Second
const CrashRestartTimeout = 30 * time.Second

//Tasks

func TaskGetter(logger lager.Logger, guid string) func() (*models.Task, error) {
	return func() (*models.Task, error) {
		return bbsClient.TaskByGuid(logger, guid)
	}
}

func TasksByDomainGetter(logger lager.Logger, domain string) func() ([]*models.Task, error) {
	return func() ([]*models.Task, error) {
		return bbsClient.TasksByDomain(logger, domain)
	}
}

func ClearOutTasksInDomain(domain string) {
	tasks, err := bbsClient.TasksByDomain(logger, domain)
	Expect(err).NotTo(HaveOccurred())
	for _, task := range tasks {
		if task.State != models.Task_Completed {
			bbsClient.CancelTask(logger, task.TaskGuid)
			Eventually(TaskGetter(logger, task.TaskGuid)).Should(HaveTaskState(models.Task_Completed))
		}
		Expect(bbsClient.ResolvingTask(logger, task.TaskGuid)).To(Succeed())
		Expect(bbsClient.DeleteTask(logger, task.TaskGuid)).To(Succeed())
	}
	Eventually(TasksByDomainGetter(logger, domain)).Should(BeEmpty())
}

func Task() *models.TaskDefinition {
	return &models.TaskDefinition{
		Action: models.WrapAction(&models.RunAction{
			Path: "bash",
			Args: []string{"-c", "echo 'some output' > /tmp/bar"},
			User: "vcap",
		}),
		RootFs:        config.DefaultRootFS,
		MemoryMb:      128,
		DiskMb:        128,
		CpuWeight:     100,
		LogGuid:       guid,
		LogSource:     "VIZ",
		ResultFile:    "/tmp/bar",
		Annotation:    "arbitrary-data",
		PlacementTags: PlacementTags(),
	}
}

//LRPs

func LRPGetter(logger lager.Logger, guid string) func() (*models.DesiredLRP, error) {
	return func() (*models.DesiredLRP, error) {
		return bbsClient.DesiredLRPByProcessGuid(logger, guid)
	}
}

func ActualLRPByProcessGuidAndIndex(logger lager.Logger, guid string, index int) (models.ActualLRP, error) {
	i := int32(index)
	lrps, err := bbsClient.ActualLRPs(logger, models.ActualLRPFilter{ProcessGuid: guid, Index: &i})
	if err != nil {
		return models.ActualLRP{}, err
	}
	if len(lrps) != 1 {
		return models.ActualLRP{}, fmt.Errorf("found more than one or no matching ActualLRP ProcessGuid: %s Index: %d", guid, index)
	}
	return *lrps[0], nil
}

func ActualsByProcessGuid(logger lager.Logger, guid string) ([]models.ActualLRP, error) {
	lrps, err := bbsClient.ActualLRPs(logger, models.ActualLRPFilter{ProcessGuid: guid})
	if err != nil {
		return nil, err
	}
	actualLRPs := make([]models.ActualLRP, len(lrps))
	for k, v := range lrps {
		actualLRPs[k] = *v
	}

	return actualLRPs, nil
}

func ActualsByDomain(logger lager.Logger, domain string) ([]models.ActualLRP, error) {
	lrps, err := bbsClient.ActualLRPs(logger, models.ActualLRPFilter{Domain: domain})
	if err != nil {
		return nil, err
	}

	actualLRPs := make([]models.ActualLRP, len(lrps))
	for k, v := range lrps {
		actualLRPs[k] = *v
	}

	return actualLRPs, nil
}

func ActualGetter(logger lager.Logger, guid string, index int) func() (models.ActualLRP, error) {
	return func() (models.ActualLRP, error) {
		return ActualLRPByProcessGuidAndIndex(logger, guid, index)
	}
}

func ActualByProcessGuidGetter(logger lager.Logger, guid string) func() ([]models.ActualLRP, error) {
	return func() ([]models.ActualLRP, error) {
		return ActualsByProcessGuid(logger, guid)
	}
}

func ActualByDomainGetter(logger lager.Logger, domain string) func() ([]models.ActualLRP, error) {
	return func() ([]models.ActualLRP, error) {
		return ActualsByDomain(logger, domain)
	}
}

func ClearOutDesiredLRPsInDomain(domain string) {
	lrps, err := bbsClient.DesiredLRPs(logger, models.DesiredLRPFilter{Domain: domain})
	Expect(err).NotTo(HaveOccurred())
	for _, lrp := range lrps {
		Expect(bbsClient.RemoveDesiredLRP(logger, lrp.ProcessGuid)).To(Succeed())
	}
	// Wait enough time for the Grace app to exit if it was run with -catchTerminate
	Eventually(ActualByDomainGetter(logger, domain), timeout+5*time.Second).Should(BeEmpty())
}

func EndpointCurler(endpoint string) func() int {
	return func() int {
		resp, err := http.Get(endpoint)
		Expect(err).NotTo(HaveOccurred())
		resp.Body.Close()
		return resp.StatusCode
	}
}

func EndpointContentCurler(endpoint string) func() string {
	return func() string {
		resp, err := http.Get(endpoint)
		Expect(err).NotTo(HaveOccurred())
		defer resp.Body.Close()
		content, err := ioutil.ReadAll(resp.Body)
		Expect(err).NotTo(HaveOccurred())
		return string(content)
	}
}

func IndexCounter(guid string, optionalHttpClient ...*http.Client) func() int {
	return IndexCounterWithAttempts(guid, 100, optionalHttpClient...)
}

func IndexCounterWithAttempts(guid string, attempts int, optionalHttpClient ...*http.Client) func() int {
	return func() int {
		counts := map[int]bool{}
		for i := 0; i < attempts; i++ {
			index := GetIndexFromEndpointFor(guid, optionalHttpClient...)
			if index == -1 {
				continue
			}
			counts[index] = true
		}
		return len(counts)
	}
}

func GetIndexFromEndpointFor(guid string, optionalHttpClient ...*http.Client) int {
	httpClient := http.DefaultClient
	if len(optionalHttpClient) == 1 {
		httpClient = optionalHttpClient[0]
	}
	url := "http://" + RouteForGuid(guid) + "/index"
	resp, err := httpClient.Get(url)
	Expect(err).NotTo(HaveOccurred())
	if resp.StatusCode != http.StatusOK {
		resp.Body.Close()
		return -1
	}
	content, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	Expect(err).NotTo(HaveOccurred())

	index, err := strconv.Atoi(string(content))
	Expect(err).NotTo(HaveOccurred())
	return index
}

func GraceCounterGetter(guid string) func() (int, error) {
	return func() (int, error) {
		resp, err := http.Get("http://" + RouteForGuid(guid) + "/counter")
		if err != nil {
			return 0, err
		}
		defer resp.Body.Close()
		content, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return 0, err
		}
		return strconv.Atoi(string(content))
	}
}

func StartedAtGetter(guid string) func() (int64, error) {
	url := "http://" + RouteForGuid(guid) + "/started-at"
	return func() (int64, error) {
		resp, err := http.Get(url)
		if err != nil {
			return 0, err
		}
		if resp.StatusCode != http.StatusOK {
			resp.Body.Close()
			return 0, errors.New(fmt.Sprintf("invalid status code: %d", resp.StatusCode))
		}
		content, err := ioutil.ReadAll(resp.Body)
		resp.Body.Close()
		if err != nil {
			return 0, err
		}
		return strconv.ParseInt(string(content), 10, 64)
	}
}

func RouteForGuid(guid string) string {
	return fmt.Sprintf("%s.%s", guid, config.RoutableDomainSuffix)
}

func DirectAddressFor(guid string, index int, containerPort uint32) string {
	actualLRP, err := ActualGetter(logger, guid, index)()
	Expect(err).NotTo(HaveOccurred())
	Expect(actualLRP).NotTo(BeZero())

	for _, portMapping := range actualLRP.Ports {
		if portMapping.ContainerPort == containerPort {
			return fmt.Sprintf("%s:%d", actualLRP.Address, portMapping.HostPort)
		}
	}

	ginkgo.Fail(fmt.Sprintf("could not find port %d for ActualLRP %d with ProcessGuid %s", containerPort, index, guid))
	return ""
}

func TLSDirectAddressFor(guid string, index int, containerPort uint32) string {
	actualLRP, err := ActualGetter(logger, guid, index)()
	Expect(err).NotTo(HaveOccurred())
	Expect(actualLRP).NotTo(BeZero())

	for _, portMapping := range actualLRP.Ports {
		if portMapping.ContainerPort == containerPort {
			return fmt.Sprintf("%s:%d", actualLRP.Address, portMapping.HostTlsProxyPort)
		}
	}

	ginkgo.Fail(fmt.Sprintf("could not find port %d for ActualLRP %d with ProcessGuid %s", containerPort, index, guid))
	return ""
}

func DesiredLRPWithGuid(guid string) *models.DesiredLRP {
	routingInfo := cfroutes.CFRoutes{
		{Port: 8080, Hostnames: []string{RouteForGuid(guid)}},
	}.RoutingInfo()

	return &models.DesiredLRP{
		ProcessGuid:   guid,
		PlacementTags: PlacementTags(),
		Domain:        domain,
		Instances:     1,
		CachedDependencies: []*models.CachedDependency{
			&models.CachedDependency{
				From:              config.GraceTarballURL,
				To:                "/tmp/grace",
				CacheKey:          "grace",
				ChecksumAlgorithm: "sha1",
				ChecksumValue:     config.GraceTarballChecksum,
			},
		},
		Action: models.WrapAction(&models.RunAction{
			Path: "/tmp/grace/grace",
			User: "vcap",
			Env:  []*models.EnvironmentVariable{{Name: "PORT", Value: "8080"}, {"ACTION_LEVEL", "COYOTE"}, {"OVERRIDE", "DAQUIRI"}},
		}),
		Monitor: models.WrapAction(&models.RunAction{
			Path: "nc",
			Args: []string{"-z", "0.0.0.0", "8080"},
			User: "vcap",
		}),
		RootFs:     config.DefaultRootFS,
		MemoryMb:   128,
		DiskMb:     128,
		CpuWeight:  100,
		Ports:      []uint32{8080},
		Routes:     &routingInfo,
		LogGuid:    guid,
		LogSource:  "VIZ",
		Annotation: "arbitrary-data",
	}
}

func PlacementTags() []string {
	return config.RepPlacementTags
}
<commit_msg>Increase timeout to make sure containers are removed<commit_after>package vizzini_test

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"
	"time"

	"code.cloudfoundry.org/bbs/models"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/routing-info/cfroutes"

	. "code.cloudfoundry.org/vizzini/matchers"
	"github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

const HealthyCheckInterval = 30 * time.Second
const ConvergerInterval = 30 * time.Second
const CrashRestartTimeout = 30 * time.Second

//Tasks

func TaskGetter(logger lager.Logger, guid string) func() (*models.Task, error) {
	return func() (*models.Task, error) {
		return bbsClient.TaskByGuid(logger, guid)
	}
}

func TasksByDomainGetter(logger lager.Logger, domain string) func() ([]*models.Task, error) {
	return func() ([]*models.Task, error) {
		return bbsClient.TasksByDomain(logger, domain)
	}
}

func ClearOutTasksInDomain(domain string) {
	tasks, err := bbsClient.TasksByDomain(logger, domain)
	Expect(err).NotTo(HaveOccurred())
	for _, task := range tasks {
		if task.State != models.Task_Completed {
			bbsClient.CancelTask(logger, task.TaskGuid)
			Eventually(TaskGetter(logger, task.TaskGuid)).Should(HaveTaskState(models.Task_Completed))
		}
		Expect(bbsClient.ResolvingTask(logger, task.TaskGuid)).To(Succeed())
		Expect(bbsClient.DeleteTask(logger, task.TaskGuid)).To(Succeed())
	}
	Eventually(TasksByDomainGetter(logger, domain)).Should(BeEmpty())
}

func Task() *models.TaskDefinition {
	return &models.TaskDefinition{
		Action: models.WrapAction(&models.RunAction{
			Path: "bash",
			Args: []string{"-c", "echo 'some output' > /tmp/bar"},
			User: "vcap",
		}),
		RootFs:        config.DefaultRootFS,
		MemoryMb:      128,
		DiskMb:        128,
		CpuWeight:     100,
		LogGuid:       guid,
		LogSource:     "VIZ",
		ResultFile:    "/tmp/bar",
		Annotation:    "arbitrary-data",
		PlacementTags: PlacementTags(),
	}
}

//LRPs

func LRPGetter(logger lager.Logger, guid string) func() (*models.DesiredLRP, error) {
	return func() (*models.DesiredLRP, error) {
		return bbsClient.DesiredLRPByProcessGuid(logger, guid)
	}
}

func ActualLRPByProcessGuidAndIndex(logger lager.Logger, guid string, index int) (models.ActualLRP, error) {
	i := int32(index)
	lrps, err := bbsClient.ActualLRPs(logger, models.ActualLRPFilter{ProcessGuid: guid, Index: &i})
	if err != nil {
		return models.ActualLRP{}, err
	}
	if len(lrps) != 1 {
		return models.ActualLRP{}, fmt.Errorf("found more than one or no matching ActualLRP ProcessGuid: %s Index: %d", guid, index)
	}
	return *lrps[0], nil
}

func ActualsByProcessGuid(logger lager.Logger, guid string) ([]models.ActualLRP, error) {
	lrps, err := bbsClient.ActualLRPs(logger, models.ActualLRPFilter{ProcessGuid: guid})
	if err != nil {
		return nil, err
	}
	actualLRPs := make([]models.ActualLRP, len(lrps))
	for k, v := range lrps {
		actualLRPs[k] = *v
	}

	return actualLRPs, nil
}

func ActualsByDomain(logger lager.Logger, domain string) ([]models.ActualLRP, error) {
	lrps, err := bbsClient.ActualLRPs(logger, models.ActualLRPFilter{Domain: domain})
	if err != nil {
		return nil, err
	}

	actualLRPs := make([]models.ActualLRP, len(lrps))
	for k, v := range lrps {
		actualLRPs[k] = *v
	}

	return actualLRPs, nil
}

func ActualGetter(logger lager.Logger, guid string, index int) func() (models.ActualLRP, error) {
	return func() (models.ActualLRP, error) {
		return ActualLRPByProcessGuidAndIndex(logger, guid, index)
	}
}

func ActualByProcessGuidGetter(logger lager.Logger, guid string) func() ([]models.ActualLRP, error) {
	return func() ([]models.ActualLRP, error) {
		return ActualsByProcessGuid(logger, guid)
	}
}

func ActualByDomainGetter(logger lager.Logger, domain string) func() ([]models.ActualLRP, error) {
	return func() ([]models.ActualLRP, error) {
		return ActualsByDomain(logger, domain)
	}
}

func ClearOutDesiredLRPsInDomain(domain string) {
	lrps, err := bbsClient.DesiredLRPs(logger, models.DesiredLRPFilter{Domain: domain})
	Expect(err).NotTo(HaveOccurred())
	for _, lrp := range lrps {
		Expect(bbsClient.RemoveDesiredLRP(logger, lrp.ProcessGuid)).To(Succeed())
	}
	// Wait enough time for the Grace app to exit if it was run with -catchTerminate
	Eventually(ActualByDomainGetter(logger, domain), timeout+8*time.Second).Should(BeEmpty())
}

func EndpointCurler(endpoint string) func() int {
	return func() int {
		resp, err := http.Get(endpoint)
		Expect(err).NotTo(HaveOccurred())
		resp.Body.Close()
		return resp.StatusCode
	}
}

func EndpointContentCurler(endpoint string) func() string {
	return func() string {
		resp, err := http.Get(endpoint)
		Expect(err).NotTo(HaveOccurred())
		defer resp.Body.Close()
		content, err := ioutil.ReadAll(resp.Body)
		Expect(err).NotTo(HaveOccurred())
		return string(content)
	}
}

func IndexCounter(guid string, optionalHttpClient ...*http.Client) func() int {
	return IndexCounterWithAttempts(guid, 100, optionalHttpClient...)
}

func IndexCounterWithAttempts(guid string, attempts int, optionalHttpClient ...*http.Client) func() int {
	return func() int {
		counts := map[int]bool{}
		for i := 0; i < attempts; i++ {
			index := GetIndexFromEndpointFor(guid, optionalHttpClient...)
			if index == -1 {
				continue
			}
			counts[index] = true
		}
		return len(counts)
	}
}

func GetIndexFromEndpointFor(guid string, optionalHttpClient ...*http.Client) int {
	httpClient := http.DefaultClient
	if len(optionalHttpClient) == 1 {
		httpClient = optionalHttpClient[0]
	}
	url := "http://" + RouteForGuid(guid) + "/index"
	resp, err := httpClient.Get(url)
	Expect(err).NotTo(HaveOccurred())
	if resp.StatusCode != http.StatusOK {
		resp.Body.Close()
		return -1
	}
	content, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	Expect(err).NotTo(HaveOccurred())

	index, err := strconv.Atoi(string(content))
	Expect(err).NotTo(HaveOccurred())
	return index
}

func GraceCounterGetter(guid string) func() (int, error) {
	return func() (int, error) {
		resp, err := http.Get("http://" + RouteForGuid(guid) + "/counter")
		if err != nil {
			return 0, err
		}
		defer resp.Body.Close()
		content, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return 0, err
		}
		return strconv.Atoi(string(content))
	}
}

func StartedAtGetter(guid string) func() (int64, error) {
	url := "http://" + RouteForGuid(guid) + "/started-at"
	return func() (int64, error) {
		resp, err := http.Get(url)
		if err != nil {
			return 0, err
		}
		if resp.StatusCode != http.StatusOK {
			resp.Body.Close()
			return 0, errors.New(fmt.Sprintf("invalid status code: %d", resp.StatusCode))
		}
		content, err := ioutil.ReadAll(resp.Body)
		resp.Body.Close()
		if err != nil {
			return 0, err
		}
		return strconv.ParseInt(string(content), 10, 64)
	}
}

func RouteForGuid(guid string) string {
	return fmt.Sprintf("%s.%s", guid, config.RoutableDomainSuffix)
}

func DirectAddressFor(guid string, index int, containerPort uint32) string {
	actualLRP, err := ActualGetter(logger, guid, index)()
	Expect(err).NotTo(HaveOccurred())
	Expect(actualLRP).NotTo(BeZero())

	for _, portMapping := range actualLRP.Ports {
		if portMapping.ContainerPort == containerPort {
			return fmt.Sprintf("%s:%d", actualLRP.Address, portMapping.HostPort)
		}
	}

	ginkgo.Fail(fmt.Sprintf("could not find port %d for ActualLRP %d with ProcessGuid %s", containerPort, index, guid))
	return ""
}

func TLSDirectAddressFor(guid string, index int, containerPort uint32) string {
	actualLRP, err := ActualGetter(logger, guid, index)()
	Expect(err).NotTo(HaveOccurred())
	Expect(actualLRP).NotTo(BeZero())

	for _, portMapping := range actualLRP.Ports {
		if portMapping.ContainerPort == containerPort {
			return fmt.Sprintf("%s:%d", actualLRP.Address, portMapping.HostTlsProxyPort)
		}
	}

	ginkgo.Fail(fmt.Sprintf("could not find port %d for ActualLRP %d with ProcessGuid %s", containerPort, index, guid))
	return ""
}

func DesiredLRPWithGuid(guid string) *models.DesiredLRP {
	routingInfo := cfroutes.CFRoutes{
		{Port: 8080, Hostnames: []string{RouteForGuid(guid)}},
	}.RoutingInfo()

	return &models.DesiredLRP{
		ProcessGuid:   guid,
		PlacementTags: PlacementTags(),
		Domain:        domain,
		Instances:     1,
		CachedDependencies: []*models.CachedDependency{
			&models.CachedDependency{
				From:              config.GraceTarballURL,
				To:                "/tmp/grace",
				CacheKey:          "grace",
				ChecksumAlgorithm: "sha1",
				ChecksumValue:     config.GraceTarballChecksum,
			},
		},
		Action: models.WrapAction(&models.RunAction{
			Path: "/tmp/grace/grace",
			User: "vcap",
			Env:  []*models.EnvironmentVariable{{Name: "PORT", Value: "8080"}, {"ACTION_LEVEL", "COYOTE"}, {"OVERRIDE", "DAQUIRI"}},
		}),
		Monitor: models.WrapAction(&models.RunAction{
			Path: "nc",
			Args: []string{"-z", "0.0.0.0", "8080"},
			User: "vcap",
		}),
		RootFs:     config.DefaultRootFS,
		MemoryMb:   128,
		DiskMb:     128,
		CpuWeight:  100,
		Ports:      []uint32{8080},
		Routes:     &routingInfo,
		LogGuid:    guid,
		LogSource:  "VIZ",
		Annotation: "arbitrary-data",
	}
}

func PlacementTags() []string {
	return config.RepPlacementTags
}
<|endoftext|>


<commit_before>package gogadgets

import (
	"labix.org/v2/mgo"
)

//Recorder takes all the update messages it receives and saves them
//in a mongodb.
type Recorder struct {
	OutputDevice
	DBHost string
	DBName string
	session *mgo.Session
	collection *mgo.Collection
	status bool
	connected bool
}

func NewRecorder(pin *Pin) (OutputDevice, error) {
	return &Recorder{
		DBHost: pin.Args["host"],
		DBName: pin.Args["db"],
	}, nil
}

func (r *Recorder) Update(msg *Message) {
	if msg.Type == "update" {
		r.save(msg)
	}
}

func (r *Recorder) On(val *Value) error {
	err := r.connect()
	r.status = true
	return err
}

func (r *Recorder) Off() error {
	r.session.Close()
	r.status = false
	return nil
}

func (r *Recorder) Status() interface{} {
	return r.status
}

func (r *Recorder)save(msg *Message) {
	r.collection.Insert(msg)
}

func (r *Recorder)connect() error {
	session, err := mgo.Dial(r.DBHost)
	if err != nil {
		return err
        }
	r.session = session
	r.collection = session.DB(r.DBName).C("updates")
	return nil
}
<commit_msg>check errors in recorder<commit_after>package gogadgets

import (
	"labix.org/v2/mgo"
)

//Recorder takes all the update messages it receives and saves them
//in a mongodb.
type Recorder struct {
	OutputDevice
	DBHost string
	DBName string
	session *mgo.Session
	collection *mgo.Collection
	status bool
	connected bool
}

func NewRecorder(pin *Pin) (OutputDevice, error) {
	return &Recorder{
		DBHost: pin.Args["host"],
		DBName: pin.Args["db"],
	}, nil
}

func (r *Recorder) Update(msg *Message) {
	if r.status && msg.Type == "update" {
		r.save(msg)
	}
}

func (r *Recorder) On(val *Value) error {
	err := r.connect()
	if err != nil {
		r.status = true
	}
	return err
}

func (r *Recorder) Off() error {
	r.session.Close()
	r.status = false
	return nil
}

func (r *Recorder) Status() interface{} {
	return r.status
}

func (r *Recorder)save(msg *Message) {
	r.collection.Insert(msg)
}

func (r *Recorder)connect() error {
	session, err := mgo.Dial(r.DBHost)
	if err != nil {
		return err
        }
	r.session = session
	r.collection = session.DB(r.DBName).C("updates")
	return nil
}
<|endoftext|>


<commit_before>package hhm

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"strconv"
	"strings"

	"github.com/TuftsBCB/io/fasta"
	"github.com/TuftsBCB/seq"
)

func Read(r io.Reader) (*HHM, error) {
	// An hhm file as four logical sections: 1) Meta data, 2) secondary
	// structure info (optional), 3) A2M formatted MSA and 4) the HMM.
	// We group 2+3 together, and store each of the three portions in their
	// own buffer. We then parse each separately.
	bmeta, bseq, bhmm := new(bytes.Buffer), new(bytes.Buffer), new(bytes.Buffer)
	mode := 1 // 1 for meta, 2 for sequence and 3 for hmm

	buf := bufio.NewReader(r)
MAIN:
	for {
		line, err := buf.ReadBytes('\n')
		if err == io.EOF && len(line) == 0 {
			break
		}
		if err != nil && err != io.EOF {
			return nil, fmt.Errorf("Error reading hhm: %s", err)
		}
		line = trim(line)

		// First check if we should do a mode change.
		// i.e., a line starting with 'SEQ' means mode 1 -> 2, and
		// a line starting with '#' means mode 2 -> 3.
		// Finally, a line starting with '\\' means STOP.
		switch {
		case hasPrefix(line, "SEQ"):
			mode = 2
			continue MAIN
		case hasPrefix(line, "#"):
			mode = 3
			continue MAIN
		case hasPrefix(line, "//"):
			break MAIN
		}

		// Now add the bytes to the appropriate buffer based on the mode.
		switch mode {
		case 1: // meta data
			_, err = bmeta.Write(append(line, '\n'))
		case 2: // sequences
			_, err = bseq.Write(append(line, '\n'))
		case 3: // hmm
			_, err = bhmm.Write(append(line, '\n'))
		default:
			panic(fmt.Sprintf("BUG: Unknown mode: %d", mode))
		}
		if err != nil {
			return nil, fmt.Errorf("Error updating buffer for hhm: %s", err)
		}
	}

	meta, err := readMeta(bmeta)
	if err != nil {
		return nil, fmt.Errorf("Error reading meta data from hhm: %s", err)
	}

	ss, msa, err := readSeqs(bseq)
	if err != nil {
		return nil, fmt.Errorf("Error reading sequence data from hhm: %s", err)
	}

	hmm, err := readHMM(bhmm)
	if err != nil {
		return nil, fmt.Errorf("Error reading HMM data from hhm: %s", err)
	}

	return &HHM{
		Meta:      meta,
		Secondary: ss,
		MSA:       msa,
		HMM:       hmm,
	}, nil
}

func readMeta(buf *bytes.Buffer) (HHMMeta, error) {
	meta := HHMMeta{}
	for {
		line, err := buf.ReadBytes('\n')
		if err == io.EOF && len(line) == 0 {
			break
		}
		if err != nil && err != io.EOF {
			return HHMMeta{}, err
		}

		line = trim(line)
		switch {
		case hasPrefix(line, "HH"):
			meta.FormatVersion = str(line)
		case hasPrefix(line, "NAME"):
			meta.Name = str(line[4:])
		case hasPrefix(line, "FAM"):
			meta.Fam = str(line[3:])
		case hasPrefix(line, "FILE"):
			meta.File = str(line[4:])
		case hasPrefix(line, "LENG"):
			meta.Leng = str(line[4:])
		case hasPrefix(line, "FILT"):
			meta.Filt = str(line[4:])
		case hasPrefix(line, "NEFF"):
			// You'd think we could use readNeff here, but does the HHM
			// format store all Neff values equally? NOOOOOOOOOOOOOOOOOOOO.
			f, err := strconv.ParseFloat(str(line[4:]), 64)
			if err != nil {
				return HHMMeta{}, err
			}
			meta.Neff = seq.Prob(f)
		case hasPrefix(line, "EVD"):
			fields := bytes.Fields(bytes.TrimSpace(line[3:]))
			if len(fields) != 2 {
				return HHMMeta{}, fmt.Errorf("Invalid EVD format: '%s'", line)
			}

			lambda, err := strconv.ParseFloat(string(fields[0]), 64)
			if err != nil {
				return HHMMeta{}, fmt.Errorf("Error EVD lambda '%s': %s",
					string(fields[0]), err)
			}
			meta.EvdLambda = lambda

			mu, err := strconv.ParseFloat(string(fields[1]), 64)
			if err != nil {
				return HHMMeta{}, fmt.Errorf("Error EVD mu '%s': %s",
					string(fields[1]), err)
			}
			meta.EvdMu = mu
		case hasPrefix(line, "PCT"):
			meta.Pct = true
		case hasPrefix(line, "DESC"):
			meta.Desc = str(line[4:])
		case hasPrefix(line, "COM"):
			meta.Com = str(line[3:])
		case hasPrefix(line, "DATE"):
			meta.Date = str(line[4:])
		}
	}
	return meta, nil
}

func readSeqs(buf *bytes.Buffer) (HHMSecondary, seq.MSA, error) {
	// Remember, the sequence portion of an HHM file actually has two parts.
	// The first part is optional and contains secondary structure information.
	// These SS sequences can be identified by special sequence headers:
	// "ss_dssp", "sa_dssp", "ss_pred", "ss_conf", and "Consensus".
	// If a sequence doesn't contain a special header, then that signifies that
	// we should start reading the MSA, which comes after the SS information.
	doneSS := false

	ss := HHMSecondary{}
	msa := seq.NewMSA()

	reader := fasta.NewReader(buf)
	reader.TrustSequences = true
	seqs, err := reader.ReadAll()
	if err != nil {
		return HHMSecondary{}, seq.MSA{}, err
	}

	for _, s := range seqs {
		s := s
		if !doneSS {
			switch {
			case strings.HasPrefix(s.Name, "ss_dssp"):
				ss.SSdssp = &s
			case strings.HasPrefix(s.Name, "sa_dssp"):
				ss.SAdssp = &s
			case strings.HasPrefix(s.Name, "ss_pred"):
				ss.SSpred = &s
			case strings.HasPrefix(s.Name, "ss_conf"):
				ss.SSconf = &s
			case strings.HasPrefix(s.Name, "Consensus"):
				ss.Consensus = &s
			default:
				doneSS = true
			}
		}
		if doneSS {
			msa.Add(s)
		}
	}
	return ss, msa, nil
}

func readHMM(buf *bytes.Buffer) (hmm *seq.HMM, err error) {
	var nullFields []string
	hmm = new(seq.HMM)
	for {
		line, err := buf.ReadBytes('\n')
		if err == io.EOF && len(line) == 0 {
			break
		}
		if err != nil && err != io.EOF {
			panic(err)
		}
		line = trim(line)

		if len(line) == 0 {
			continue
		}
		switch {
		case hasPrefix(line, "NULL"):
			// We can't read the NULL emissions yet, because we don't have
			// an alphabet. (Which we'll get on the next line.)
			// We'll slurp this into a seq.EProbs value in a little bit, after
			// we get an alphabet.
			nullFields = strings.Fields(str(line[4:]))
		case hasPrefix(line, "HMM"):
			// We slurp up three lines here. The first is the alphabet
			// (the current line). The second is the ordering of transition
			// probabilities. And the third are transition probabilities for
			// the begin state. We ignore the second two.
			if _, err := demandLine(buf); err != nil {
				return nil, fmt.Errorf("%s (expected transition ordering)", err)
			}
			if _, err := demandLine(buf); err != nil {
				return nil, fmt.Errorf("%s (expected start transitions)", err)
			}

			// Get the ordering of the alphabet.
			hmm.Alphabet = make([]seq.Residue, 0, 20)
			residues := bytes.Split(trim(line[3:]), []byte{'\t'})
			for _, residue := range residues {
				hmm.Alphabet = append(hmm.Alphabet, seq.Residue(residue[0]))
			}

			// Remember those null probabilities? Well, we have an alphabet now.
			ep, err := readEmissions(hmm.Alphabet, nullFields)
			if err != nil {
				return nil, fmt.Errorf("Could not read NULL emissions '%s': %s",
					strings.Join(nullFields, " "), err)
			}
			hmm.Null = ep
		default: // finally, reading a node in the HMM
			// Each node in the HMM is made up of two lines.
			// The first line starts with the amino acid in the reference
			// sequence, followed by the node number, followed by N match
			// emission probabilities (where N = len(alphabet)), and finally
			// followed by the node number again. (What the *fuck* is up with
			// that? Seriously.)
			//
			// The second line is made up of 7 transition probabilities,
			// followed by 3 diversity (the 'neff' stuff) scores.
			//
			// Also, each field is separated by spaces OR tabs. Lovely, eh?
			line2, err := demandLine(buf)
			if err != nil {
				return nil, fmt.Errorf("%s (expected transition probs)", err)
			}
			fields1 := strings.Fields(string(line))
			fields2 := strings.Fields(string(line2))
			node := seq.HMMNode{
				HMM:     hmm,
				Residue: seq.Residue(fields1[0][0]),
			}

			node.NodeNum, err = strconv.Atoi(fields1[1])
			if err != nil {
				return nil, fmt.Errorf("Could not parse node number '%s': %s",
					fields1[1], err)
			}

			node.MatEmit, err = readEmissions(hmm.Alphabet, fields1[2:])
			if err != nil {
				return nil, fmt.Errorf("Could not read emissions '%s': %s",
					strings.Join(fields1[2:], " "), err)
			}

			node.InsEmit = seq.NewEProbs(hmm.Alphabet)
			for residue := range node.InsEmit {
				node.InsEmit[residue] = hmm.Null[residue]
			}

			node.Transitions, err = readTransitions(fields2)
			if err != nil {
				return nil, fmt.Errorf("Could not read transitions '%s': %s",
					strings.Join(fields2, " "), err)
			}

			node.NeffM, node.NeffI, node.NeffD, err = readDiversity(fields2[7:])
			if err != nil {
				return nil, fmt.Errorf("Could not read diversity '%s': %s",
					strings.Join(fields2[7:], " "), err)
			}

			hmm.Nodes = append(hmm.Nodes, node)
		}
	}
	return hmm, nil
}

func readEmissions(alphabet []seq.Residue, flds []string) (seq.EProbs, error) {
	var err error

	ep := seq.NewEProbs(alphabet)
	for i := 0; i < len(alphabet); i++ {
		if ep[alphabet[i]], err = readProb(flds[i]); err != nil {
			return nil, err
		}
	}
	return ep, nil
}

func readTransitions(fields []string) (tp seq.TProbs, err error) {
	if tp.MM, err = readProb(fields[0]); err != nil {
		return
	}
	if tp.MI, err = readProb(fields[1]); err != nil {
		return
	}
	if tp.MD, err = readProb(fields[2]); err != nil {
		return
	}
	if tp.IM, err = readProb(fields[3]); err != nil {
		return
	}
	if tp.II, err = readProb(fields[4]); err != nil {
		return
	}
	if tp.DM, err = readProb(fields[5]); err != nil {
		return
	}
	if tp.DD, err = readProb(fields[6]); err != nil {
		return
	}
	return
}

func readDiversity(fields []string) (m, i, d seq.Prob, err error) {
	if m, err = readNeff(fields[0]); err != nil {
		return
	}
	if i, err = readNeff(fields[1]); err != nil {
		return
	}
	if d, err = readNeff(fields[2]); err != nil {
		return
	}
	return
}

func demandLine(buf *bytes.Buffer) ([]byte, error) {
	line, err := buf.ReadBytes('\n')
	if err == io.EOF {
		return nil, fmt.Errorf("Unexpected EOF when reading HMM in hhm.")
	}
	if err != nil {
		return nil, err
	}
	return trim(line), nil
}

// readProb reads a probability (transition or emissions) from an hhm file and
// returns a Prob value in log_2 form.
func readProb(fstr string) (seq.Prob, error) {
	f, err := seq.NewProb(fstr)
	if err != nil {
		return f, fmt.Errorf("Error reading probability '%s': %s", fstr, err)
	}
	if f.IsMin() {
		return f, nil
	}
	return -f / seq.Prob(hmmScale), nil
}

// readNeff reads a diversity value.
func readNeff(fstr string) (seq.Prob, error) {
	f, err := seq.NewProb(fstr)
	if err != nil {
		return f, fmt.Errorf("Error reading neff '%s': %s", fstr, err)
	}
	if f.IsMin() {
		return f, nil
	}
	return f / seq.Prob(hmmScale), nil
}

func hasPrefix(bs []byte, prefix string) bool {
	return bytes.HasPrefix(bs, []byte(prefix))
}

func trim(bs []byte) []byte {
	return bytes.TrimSpace(bs)
}

func str(bs []byte) string {
	return string(bytes.TrimSpace(bs))
}
<commit_msg>HMM nodes no longer refer back to HMM.<commit_after>package hhm

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"strconv"
	"strings"

	"github.com/TuftsBCB/io/fasta"
	"github.com/TuftsBCB/seq"
)

func Read(r io.Reader) (*HHM, error) {
	// An hhm file as four logical sections: 1) Meta data, 2) secondary
	// structure info (optional), 3) A2M formatted MSA and 4) the HMM.
	// We group 2+3 together, and store each of the three portions in their
	// own buffer. We then parse each separately.
	bmeta, bseq, bhmm := new(bytes.Buffer), new(bytes.Buffer), new(bytes.Buffer)
	mode := 1 // 1 for meta, 2 for sequence and 3 for hmm

	buf := bufio.NewReader(r)
MAIN:
	for {
		line, err := buf.ReadBytes('\n')
		if err == io.EOF && len(line) == 0 {
			break
		}
		if err != nil && err != io.EOF {
			return nil, fmt.Errorf("Error reading hhm: %s", err)
		}
		line = trim(line)

		// First check if we should do a mode change.
		// i.e., a line starting with 'SEQ' means mode 1 -> 2, and
		// a line starting with '#' means mode 2 -> 3.
		// Finally, a line starting with '\\' means STOP.
		switch {
		case hasPrefix(line, "SEQ"):
			mode = 2
			continue MAIN
		case hasPrefix(line, "#"):
			mode = 3
			continue MAIN
		case hasPrefix(line, "//"):
			break MAIN
		}

		// Now add the bytes to the appropriate buffer based on the mode.
		switch mode {
		case 1: // meta data
			_, err = bmeta.Write(append(line, '\n'))
		case 2: // sequences
			_, err = bseq.Write(append(line, '\n'))
		case 3: // hmm
			_, err = bhmm.Write(append(line, '\n'))
		default:
			panic(fmt.Sprintf("BUG: Unknown mode: %d", mode))
		}
		if err != nil {
			return nil, fmt.Errorf("Error updating buffer for hhm: %s", err)
		}
	}

	meta, err := readMeta(bmeta)
	if err != nil {
		return nil, fmt.Errorf("Error reading meta data from hhm: %s", err)
	}

	ss, msa, err := readSeqs(bseq)
	if err != nil {
		return nil, fmt.Errorf("Error reading sequence data from hhm: %s", err)
	}

	hmm, err := readHMM(bhmm)
	if err != nil {
		return nil, fmt.Errorf("Error reading HMM data from hhm: %s", err)
	}

	return &HHM{
		Meta:      meta,
		Secondary: ss,
		MSA:       msa,
		HMM:       hmm,
	}, nil
}

func readMeta(buf *bytes.Buffer) (HHMMeta, error) {
	meta := HHMMeta{}
	for {
		line, err := buf.ReadBytes('\n')
		if err == io.EOF && len(line) == 0 {
			break
		}
		if err != nil && err != io.EOF {
			return HHMMeta{}, err
		}

		line = trim(line)
		switch {
		case hasPrefix(line, "HH"):
			meta.FormatVersion = str(line)
		case hasPrefix(line, "NAME"):
			meta.Name = str(line[4:])
		case hasPrefix(line, "FAM"):
			meta.Fam = str(line[3:])
		case hasPrefix(line, "FILE"):
			meta.File = str(line[4:])
		case hasPrefix(line, "LENG"):
			meta.Leng = str(line[4:])
		case hasPrefix(line, "FILT"):
			meta.Filt = str(line[4:])
		case hasPrefix(line, "NEFF"):
			// You'd think we could use readNeff here, but does the HHM
			// format store all Neff values equally? NOOOOOOOOOOOOOOOOOOOO.
			f, err := strconv.ParseFloat(str(line[4:]), 64)
			if err != nil {
				return HHMMeta{}, err
			}
			meta.Neff = seq.Prob(f)
		case hasPrefix(line, "EVD"):
			fields := bytes.Fields(bytes.TrimSpace(line[3:]))
			if len(fields) != 2 {
				return HHMMeta{}, fmt.Errorf("Invalid EVD format: '%s'", line)
			}

			lambda, err := strconv.ParseFloat(string(fields[0]), 64)
			if err != nil {
				return HHMMeta{}, fmt.Errorf("Error EVD lambda '%s': %s",
					string(fields[0]), err)
			}
			meta.EvdLambda = lambda

			mu, err := strconv.ParseFloat(string(fields[1]), 64)
			if err != nil {
				return HHMMeta{}, fmt.Errorf("Error EVD mu '%s': %s",
					string(fields[1]), err)
			}
			meta.EvdMu = mu
		case hasPrefix(line, "PCT"):
			meta.Pct = true
		case hasPrefix(line, "DESC"):
			meta.Desc = str(line[4:])
		case hasPrefix(line, "COM"):
			meta.Com = str(line[3:])
		case hasPrefix(line, "DATE"):
			meta.Date = str(line[4:])
		}
	}
	return meta, nil
}

func readSeqs(buf *bytes.Buffer) (HHMSecondary, seq.MSA, error) {
	// Remember, the sequence portion of an HHM file actually has two parts.
	// The first part is optional and contains secondary structure information.
	// These SS sequences can be identified by special sequence headers:
	// "ss_dssp", "sa_dssp", "ss_pred", "ss_conf", and "Consensus".
	// If a sequence doesn't contain a special header, then that signifies that
	// we should start reading the MSA, which comes after the SS information.
	doneSS := false

	ss := HHMSecondary{}
	msa := seq.NewMSA()

	reader := fasta.NewReader(buf)
	reader.TrustSequences = true
	seqs, err := reader.ReadAll()
	if err != nil {
		return HHMSecondary{}, seq.MSA{}, err
	}

	for _, s := range seqs {
		s := s
		if !doneSS {
			switch {
			case strings.HasPrefix(s.Name, "ss_dssp"):
				ss.SSdssp = &s
			case strings.HasPrefix(s.Name, "sa_dssp"):
				ss.SAdssp = &s
			case strings.HasPrefix(s.Name, "ss_pred"):
				ss.SSpred = &s
			case strings.HasPrefix(s.Name, "ss_conf"):
				ss.SSconf = &s
			case strings.HasPrefix(s.Name, "Consensus"):
				ss.Consensus = &s
			default:
				doneSS = true
			}
		}
		if doneSS {
			msa.Add(s)
		}
	}
	return ss, msa, nil
}

func readHMM(buf *bytes.Buffer) (hmm *seq.HMM, err error) {
	var nullFields []string
	hmm = new(seq.HMM)
	for {
		line, err := buf.ReadBytes('\n')
		if err == io.EOF && len(line) == 0 {
			break
		}
		if err != nil && err != io.EOF {
			panic(err)
		}
		line = trim(line)

		if len(line) == 0 {
			continue
		}
		switch {
		case hasPrefix(line, "NULL"):
			// We can't read the NULL emissions yet, because we don't have
			// an alphabet. (Which we'll get on the next line.)
			// We'll slurp this into a seq.EProbs value in a little bit, after
			// we get an alphabet.
			nullFields = strings.Fields(str(line[4:]))
		case hasPrefix(line, "HMM"):
			// We slurp up three lines here. The first is the alphabet
			// (the current line). The second is the ordering of transition
			// probabilities. And the third are transition probabilities for
			// the begin state. We ignore the second two.
			if _, err := demandLine(buf); err != nil {
				return nil, fmt.Errorf("%s (expected transition ordering)", err)
			}
			if _, err := demandLine(buf); err != nil {
				return nil, fmt.Errorf("%s (expected start transitions)", err)
			}

			// Get the ordering of the alphabet.
			hmm.Alphabet = make([]seq.Residue, 0, 20)
			residues := bytes.Split(trim(line[3:]), []byte{'\t'})
			for _, residue := range residues {
				hmm.Alphabet = append(hmm.Alphabet, seq.Residue(residue[0]))
			}

			// Remember those null probabilities? Well, we have an alphabet now.
			ep, err := readEmissions(hmm.Alphabet, nullFields)
			if err != nil {
				return nil, fmt.Errorf("Could not read NULL emissions '%s': %s",
					strings.Join(nullFields, " "), err)
			}
			hmm.Null = ep
		default: // finally, reading a node in the HMM
			// Each node in the HMM is made up of two lines.
			// The first line starts with the amino acid in the reference
			// sequence, followed by the node number, followed by N match
			// emission probabilities (where N = len(alphabet)), and finally
			// followed by the node number again. (What the *fuck* is up with
			// that? Seriously.)
			//
			// The second line is made up of 7 transition probabilities,
			// followed by 3 diversity (the 'neff' stuff) scores.
			//
			// Also, each field is separated by spaces OR tabs. Lovely, eh?
			line2, err := demandLine(buf)
			if err != nil {
				return nil, fmt.Errorf("%s (expected transition probs)", err)
			}
			fields1 := strings.Fields(string(line))
			fields2 := strings.Fields(string(line2))
			node := seq.HMMNode{
				Residue: seq.Residue(fields1[0][0]),
			}

			node.NodeNum, err = strconv.Atoi(fields1[1])
			if err != nil {
				return nil, fmt.Errorf("Could not parse node number '%s': %s",
					fields1[1], err)
			}

			node.MatEmit, err = readEmissions(hmm.Alphabet, fields1[2:])
			if err != nil {
				return nil, fmt.Errorf("Could not read emissions '%s': %s",
					strings.Join(fields1[2:], " "), err)
			}

			node.InsEmit = seq.NewEProbs(hmm.Alphabet)
			for residue := range node.InsEmit {
				node.InsEmit[residue] = hmm.Null[residue]
			}

			node.Transitions, err = readTransitions(fields2)
			if err != nil {
				return nil, fmt.Errorf("Could not read transitions '%s': %s",
					strings.Join(fields2, " "), err)
			}

			node.NeffM, node.NeffI, node.NeffD, err = readDiversity(fields2[7:])
			if err != nil {
				return nil, fmt.Errorf("Could not read diversity '%s': %s",
					strings.Join(fields2[7:], " "), err)
			}

			hmm.Nodes = append(hmm.Nodes, node)
		}
	}
	return hmm, nil
}

func readEmissions(alphabet []seq.Residue, flds []string) (seq.EProbs, error) {
	var err error

	ep := seq.NewEProbs(alphabet)
	for i := 0; i < len(alphabet); i++ {
		if ep[alphabet[i]], err = readProb(flds[i]); err != nil {
			return nil, err
		}
	}
	return ep, nil
}

func readTransitions(fields []string) (tp seq.TProbs, err error) {
	if tp.MM, err = readProb(fields[0]); err != nil {
		return
	}
	if tp.MI, err = readProb(fields[1]); err != nil {
		return
	}
	if tp.MD, err = readProb(fields[2]); err != nil {
		return
	}
	if tp.IM, err = readProb(fields[3]); err != nil {
		return
	}
	if tp.II, err = readProb(fields[4]); err != nil {
		return
	}
	if tp.DM, err = readProb(fields[5]); err != nil {
		return
	}
	if tp.DD, err = readProb(fields[6]); err != nil {
		return
	}
	return
}

func readDiversity(fields []string) (m, i, d seq.Prob, err error) {
	if m, err = readNeff(fields[0]); err != nil {
		return
	}
	if i, err = readNeff(fields[1]); err != nil {
		return
	}
	if d, err = readNeff(fields[2]); err != nil {
		return
	}
	return
}

func demandLine(buf *bytes.Buffer) ([]byte, error) {
	line, err := buf.ReadBytes('\n')
	if err == io.EOF {
		return nil, fmt.Errorf("Unexpected EOF when reading HMM in hhm.")
	}
	if err != nil {
		return nil, err
	}
	return trim(line), nil
}

// readProb reads a probability (transition or emissions) from an hhm file and
// returns a Prob value in log_2 form.
func readProb(fstr string) (seq.Prob, error) {
	f, err := seq.NewProb(fstr)
	if err != nil {
		return f, fmt.Errorf("Error reading probability '%s': %s", fstr, err)
	}
	if f.IsMin() {
		return f, nil
	}
	return -f / seq.Prob(hmmScale), nil
}

// readNeff reads a diversity value.
func readNeff(fstr string) (seq.Prob, error) {
	f, err := seq.NewProb(fstr)
	if err != nil {
		return f, fmt.Errorf("Error reading neff '%s': %s", fstr, err)
	}
	if f.IsMin() {
		return f, nil
	}
	return f / seq.Prob(hmmScale), nil
}

func hasPrefix(bs []byte, prefix string) bool {
	return bytes.HasPrefix(bs, []byte(prefix))
}

func trim(bs []byte) []byte {
	return bytes.TrimSpace(bs)
}

func str(bs []byte) string {
	return string(bytes.TrimSpace(bs))
}
<|endoftext|>


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

package main

import (
	"flag"
	"net"
	"os"
	"time"

	"github.com/banzaicloud/bank-vaults/operator/pkg/apis"
	"github.com/banzaicloud/bank-vaults/operator/pkg/controller"
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp"
	"k8s.io/client-go/rest"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client/config"
	"sigs.k8s.io/controller-runtime/pkg/healthz"
	"sigs.k8s.io/controller-runtime/pkg/log/zap"
	"sigs.k8s.io/controller-runtime/pkg/manager"
)

var log = ctrl.Log.WithName("cmd")

const (
	operatorNamespace      = "OPERATOR_NAMESPACE"
	watchNamespaceEnvVar   = "WATCH_NAMESPACE"
	healthProbeBindAddress = ":8080"
	metricsBindAddress     = ":8383"
)

func main() {

	syncPeriod := flag.Duration("sync_period", 30*time.Second, "SyncPeriod determines the minimum frequency at which watched resources are reconciled")
	verbose := flag.Bool("verbose", false, "enable verbose logging")

	flag.Parse()

	// The logger instantiated here can be changed to any logger
	// implementing the logr.Logger interface. This logger will
	// be propagated through the whole operator, generating
	// uniform and structured logs.
	ctrl.SetLogger(zap.New(zap.UseDevMode(*verbose)))

	var namespace string
	var err error
	namespace, isSet := os.LookupEnv(operatorNamespace)

	if !isSet {
		namespace, isSet = os.LookupEnv(watchNamespaceEnvVar)
		if !isSet {
			log.Info("No watched namespace found, watching the entire cluster")
			namespace = ""
		}
	}
	log.Info("Watched namespace: " + namespace)

	// Get a config to talk to the apiserver
	k8sConfig, err := config.GetConfig()
	if err != nil {
		log.Error(err, "")
		os.Exit(1)
	}

	leaderElectionNamespace := ""
	if !isInClusterConfig(k8sConfig) {
		leaderElectionNamespace = "default"
	}

	// Create a new Cmd to provide shared dependencies and start components
	mgr, err := manager.New(k8sConfig, manager.Options{
		Namespace:               namespace,
		LeaderElection:          true,
		LeaderElectionNamespace: leaderElectionNamespace,
		LeaderElectionID:        "vault-operator-lock",
		SyncPeriod:              syncPeriod,
		HealthProbeBindAddress:  healthProbeBindAddress,
		LivenessEndpointName:    "/",      // For Chart backwards compatibility
		ReadinessEndpointName:   "/ready", // For Chart backwards compatibility
		MetricsBindAddress:      metricsBindAddress,
	})
	if err != nil {
		log.Error(err, "")
		os.Exit(1)
	}

	err = mgr.AddReadyzCheck("ping", healthz.Ping)
	if err != nil {
		log.Error(err, "")
		os.Exit(1)
	}
	err = mgr.AddHealthzCheck("ping", healthz.Ping)
	if err != nil {
		log.Error(err, "")
		os.Exit(1)
	}

	log.Info("Registering Components.")

	// Setup Scheme for all resources
	if err := apis.AddToScheme(mgr.GetScheme()); err != nil {
		log.Error(err, "")
		os.Exit(1)
	}

	// Setup all Controllers
	if err := controller.AddToManager(mgr); err != nil {
		log.Error(err, "")
		os.Exit(1)
	}

	log.Info("Starting the Cmd.")

	// Start the Cmd
	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
		log.Error(err, "manager exited non-zero")
		os.Exit(1)
	}
}

func isInClusterConfig(k8sConfig *rest.Config) bool {
	host, port := os.Getenv("KUBERNETES_SERVICE_HOST"), os.Getenv("KUBERNETES_SERVICE_PORT")
	return k8sConfig.Host == "https://"+net.JoinHostPort(host, port)
}
<commit_msg>Update main.go (#1462)<commit_after>// Copyright © 2019 Banzai Cloud
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"net"
	"os"
	"time"

	"github.com/banzaicloud/bank-vaults/operator/pkg/apis"
	"github.com/banzaicloud/bank-vaults/operator/pkg/controller"
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp"
	"k8s.io/client-go/rest"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client/config"
	"sigs.k8s.io/controller-runtime/pkg/healthz"
	"sigs.k8s.io/controller-runtime/pkg/log/zap"
	"sigs.k8s.io/controller-runtime/pkg/manager"
)

var log = ctrl.Log.WithName("cmd")

const (
	operatorNamespace      = "OPERATOR_NAMESPACE"
	watchNamespaceEnvVar   = "WATCH_NAMESPACE"
	healthProbeBindAddress = ":8080"
	metricsBindAddress     = ":8383"
)

func main() {

	syncPeriod := flag.Duration("sync_period", 30*time.Second, "SyncPeriod determines the minimum frequency at which watched resources are reconciled")
	verbose := flag.Bool("verbose", false, "enable verbose logging")

	flag.Parse()

	// The logger instantiated here can be changed to any logger
	// implementing the logr.Logger interface. This logger will
	// be propagated through the whole operator, generating
	// uniform and structured logs.
	ctrl.SetLogger(zap.New(zap.UseDevMode(*verbose)))

	var namespace string
	var err error
	namespace, isSet := os.LookupEnv(operatorNamespace)

	if !isSet {
		namespace, isSet = os.LookupEnv(watchNamespaceEnvVar)
		if !isSet {
			log.Info("No watched namespace found, watching the entire cluster")
			namespace = ""
		}
	}
	log.Info("Watched namespace: " + namespace)

	// Get a config to talk to the apiserver
	k8sConfig, err := config.GetConfig()
	if err != nil {
		log.Error(err, "Unable to get k8s config")
		os.Exit(1)
	}

	leaderElectionNamespace := ""
	if !isInClusterConfig(k8sConfig) {
		leaderElectionNamespace = "default"
	}

	// Create a new Cmd to provide shared dependencies and start components
	mgr, err := manager.New(k8sConfig, manager.Options{
		Namespace:               namespace,
		LeaderElection:          true,
		LeaderElectionNamespace: leaderElectionNamespace,
		LeaderElectionID:        "vault-operator-lock",
		SyncPeriod:              syncPeriod,
		HealthProbeBindAddress:  healthProbeBindAddress,
		LivenessEndpointName:    "/",      // For Chart backwards compatibility
		ReadinessEndpointName:   "/ready", // For Chart backwards compatibility
		MetricsBindAddress:      metricsBindAddress,
	})
	if err != nil {
		log.Error(err, "Unable to create manager as defined")
		os.Exit(1)
	}

	err = mgr.AddReadyzCheck("ping", healthz.Ping)
	if err != nil {
		log.Error(err, "Add Readyz Check failed")
		os.Exit(1)
	}
	err = mgr.AddHealthzCheck("ping", healthz.Ping)
	if err != nil {
		log.Error(err, "Unable to add heatlh check")
		os.Exit(1)
	}

	log.Info("Registering Components.")

	// Setup Scheme for all resources
	if err := apis.AddToScheme(mgr.GetScheme()); err != nil {
		log.Error(err, "Failed to use api to add scheme")
		os.Exit(1)
	}

	// Setup all Controllers
	if err := controller.AddToManager(mgr); err != nil {
		log.Error(err, "Unable to add manager to controller")
		os.Exit(1)
	}

	log.Info("Starting the Cmd.")

	// Start the Cmd
	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
		log.Error(err, "manager exited non-zero")
		os.Exit(1)
	}
}

func isInClusterConfig(k8sConfig *rest.Config) bool {
	host, port := os.Getenv("KUBERNETES_SERVICE_HOST"), os.Getenv("KUBERNETES_SERVICE_PORT")
	return k8sConfig.Host == "https://"+net.JoinHostPort(host, port)
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/asaskevich/govalidator"
	"github.com/gin-gonic/gin"
	"zxq.co/ripple/rippleapi/common"
	"zxq.co/ripple/schiavolib"
)

func register(c *gin.Context) {
	if getContext(c).User.ID != 0 {
		resp403(c)
		return
	}
	if c.Query("stopsign") != "1" {
		u, _ := tryBotnets(c)
		if u != "" {
			simple(c, getSimpleByFilename("register/elmo.html"), nil, map[string]interface{}{
				"Username": u,
			})
			return
		}
	}
	registerResp(c)
}

func registerSubmit(c *gin.Context) {
	if getContext(c).User.ID != 0 {
		resp403(c)
		return
	}
	// check registrations are enabled
	if !registrationsEnabled() {
		registerResp(c, errorMessage{T(c, "Sorry, it's not possible to register at the moment. Please try again later.")})
		return
	}

	// check username is valid by our criteria
	username := strings.TrimSpace(c.PostForm("username"))
	if !usernameRegex.MatchString(username) {
		registerResp(c, errorMessage{T(c, "Your username must contain alphanumerical characters, spaces, or any of _[]-")})
		return
	}

	// check whether an username is e.g. cookiezi, shigetora, peppy, wubwoofwolf, loctav
	if in(strings.ToLower(username), forbiddenUsernames) {
		registerResp(c, errorMessage{T(c, "You're not allowed to register with that username.")})
		return
	}

	// check email is valid
	if !govalidator.IsEmail(c.PostForm("email")) {
		registerResp(c, errorMessage{T(c, "Please pass a valid email address.")})
		return
	}

	// passwords check (too short/too common)
	if x := validatePassword(c.PostForm("password")); x != "" {
		registerResp(c, errorMessage{T(c, x)})
		return
	}

	// usernames with both _ and spaces are not allowed
	if strings.Contains(username, "_") && strings.Contains(username, " ") {
		registerResp(c, errorMessage{T(c, "An username can't contain both underscores and spaces.")})
		return
	}

	// check whether username already exists
	if db.QueryRow("SELECT 1 FROM users WHERE username_safe = ?", safeUsername(username)).
		Scan(new(int)) != sql.ErrNoRows {
		registerResp(c, errorMessage{T(c, "An user with that username already exists!")})
		return
	}

	// check whether an user with that email already exists
	if db.QueryRow("SELECT 1 FROM users WHERE email = ?", c.PostForm("email")).
		Scan(new(int)) != sql.ErrNoRows {
		registerResp(c, errorMessage{T(c, "An user with that email address already exists!")})
		return
	}

	// recaptcha verify
	if config.RecaptchaPrivate != "" && !recaptchaCheck(c) {
		registerResp(c, errorMessage{T(c, "Captcha is invalid.")})
		return
	}

	uMulti, criteria := tryBotnets(c)
	if criteria != "" {
		schiavo.CMs.Send(
			fmt.Sprintf(
				"User **%s** registered with the same %s as %s (%s/u/%s). **POSSIBLE MULTIACCOUNT!!!**. Waiting for ingame verification...",
				username, criteria, uMulti, config.BaseURL, url.QueryEscape(uMulti),
			),
		)
	}

	// The actual registration.
	pass, err := generatePassword(c.PostForm("password"))
	if err != nil {
		resp500(c)
		return
	}

	res, err := db.Exec(`INSERT INTO users(username, username_safe, password_md5, salt, email, register_datetime, privileges, password_version)
							  VALUES (?,        ?,             ?,            '',   ?,     ?,                 ?,          2);`,
		username, safeUsername(username), pass, c.PostForm("email"), time.Now().Unix(), common.UserPrivilegePendingVerification)
	if err != nil {
		registerResp(c, errorMessage{T(c, "Whoops, an error slipped in. You might have been registered, though. I don't know.")})
		return
	}
	lid, _ := res.LastInsertId()

	db.Exec("INSERT INTO `users_stats`(id, username, user_color, user_style, ranked_score_std, playcount_std, total_score_std, ranked_score_taiko, playcount_taiko, total_score_taiko, ranked_score_ctb, playcount_ctb, total_score_ctb, ranked_score_mania, playcount_mania, total_score_mania) VALUES (?, ?, 'black', '', 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);", lid, username)

	schiavo.CMs.Send(fmt.Sprintf("User (**%s** | %s) registered from %s", username, c.PostForm("email"), clientIP(c)))

	setYCookie(int(lid), c)
	logIP(c, int(lid))

	rd.Incr("ripple:registered_users")

	addMessage(c, successMessage{T(c, "You have been successfully registered on Ripple! You now need to verify your account.")})
	getSession(c).Save()
	c.Redirect(302, "/register/verify?u="+strconv.Itoa(int(lid)))
}

func registerResp(c *gin.Context, messages ...message) {
	resp(c, 200, "register/register.html", &baseTemplateData{
		TitleBar:  "Register",
		KyutGrill: "register.jpg",
		Scripts:   []string{"https://www.google.com/recaptcha/api.js"},
		Messages:  messages,
		FormData:  normaliseURLValues(c.Request.PostForm),
	})
}

func registrationsEnabled() bool {
	var enabled bool
	db.QueryRow("SELECT value_int FROM system_settings WHERE name = 'registrations_enabled'").Scan(&enabled)
	return enabled
}

func verifyAccount(c *gin.Context) {
	if getContext(c).User.ID != 0 {
		resp403(c)
		return
	}

	i, ret := checkUInQS(c)
	if ret {
		return
	}

	sess := getSession(c)
	var rPrivileges uint64
	db.Get(&rPrivileges, "SELECT privileges FROM users WHERE id = ?", i)
	if common.UserPrivileges(rPrivileges)&common.UserPrivilegePendingVerification == 0 {
		addMessage(c, warningMessage{T(c, "Nope.")})
		sess.Save()
		c.Redirect(302, "/")
		return
	}

	resp(c, 200, "register/verify.html", &baseTemplateData{
		TitleBar:       "Verify account",
		HeadingOnRight: true,
		KyutGrill:      "welcome.jpg",
	})
}

func welcome(c *gin.Context) {
	if getContext(c).User.ID != 0 {
		resp403(c)
		return
	}

	i, ret := checkUInQS(c)
	if ret {
		return
	}

	var rPrivileges uint64
	db.Get(&rPrivileges, "SELECT privileges FROM users WHERE id = ?", i)
	if common.UserPrivileges(rPrivileges)&common.UserPrivilegePendingVerification > 0 {
		c.Redirect(302, "/register/verify?u="+c.Query("u"))
		return
	}

	t := T(c, "Welcome!")
	if common.UserPrivileges(rPrivileges)&common.UserPrivilegeNormal == 0 {
		// if the user has no UserNormal, it means they're banned = they multiaccounted
		t = T(c, "Welcome back!")
	}

	resp(c, 200, "register/welcome.html", &baseTemplateData{
		TitleBar:       t,
		HeadingOnRight: true,
		KyutGrill:      "welcome.jpg",
	})
}

// Check User In Query Is Same As User In Y Cookie
func checkUInQS(c *gin.Context) (int, bool) {
	sess := getSession(c)

	i, _ := strconv.Atoi(c.Query("u"))
	y, _ := c.Cookie("y")
	err := db.QueryRow("SELECT 1 FROM identity_tokens WHERE token = ? AND userid = ?", y, i).Scan(new(int))
	if err == sql.ErrNoRows {
		addMessage(c, warningMessage{T(c, "Nope.")})
		sess.Save()
		c.Redirect(302, "/")
		return 0, true
	}
	return i, false
}

func tryBotnets(c *gin.Context) (string, string) {
	var username string

	err := db.QueryRow("SELECT u.username FROM ip_user i INNER JOIN users u ON u.id = i.userid WHERE i.ip = ?", clientIP(c)).Scan(&username)
	if err != nil {
		if err != sql.ErrNoRows {
			c.Error(err)
		}
		return "", ""
	}
	if username != "" {
		return username, "IP"
	}

	cook, _ := c.Cookie("y")
	err = db.QueryRow("SELECT u.username FROM identity_tokens i INNER JOIN users u ON u.id = i.userid WHERE i.token = ?",
		cook).Scan(&username)
	if err != nil {
		if err != sql.ErrNoRows {
			c.Error(err)
		}
		return "", ""
	}
	if username != "" {
		return username, "username"
	}

	return "", ""
}

func in(s string, ss []string) bool {
	for _, x := range ss {
		if x == s {
			return true
		}
	}
	return false
}

var usernameRegex = regexp.MustCompile(`^[A-Za-z0-9 _\[\]-]{2,15}$`)
var forbiddenUsernames = []string{
	"peppy",
	"rrtyui",
	"cookiezi",
	"azer",
	"loctav",
	"banchobot",
	"happystick",
	"doomsday",
	"sharingan33",
	"andrea",
	"cptnxn",
	"reimu-desu",
	"hvick225",
	"_index",
	"my aim sucks",
	"kynan",
	"rafis",
	"sayonara-bye",
	"thelewa",
	"wubwoofwolf",
	"millhioref",
	"tom94",
	"tillerino",
	"clsw",
	"spectator",
	"exgon",
	"axarious",
	"angelsim",
	"recia",
	"nara",
	"emperorpenguin83",
	"bikko",
	"xilver",
	"vettel",
	"kuu01",
	"_yu68",
	"tasuke912",
	"dusk",
	"ttobas",
	"velperk",
	"jakads",
	"jhlee0133",
	"abcdullah",
	"yuko-",
	"entozer",
	"hdhr",
	"ekoro",
	"snowwhite",
	"osuplayer111",
	"musty",
	"nero",
	"elysion",
	"ztrot",
	"koreapenguin",
	"fort",
	"asphyxia",
	"niko",
	"shigetora",
}
<commit_msg>Populate users_stats_relax when registering<commit_after>package main

import (
	"database/sql"
	"fmt"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/asaskevich/govalidator"
	"github.com/gin-gonic/gin"
	"zxq.co/ripple/rippleapi/common"
	schiavo "zxq.co/ripple/schiavolib"
)

func register(c *gin.Context) {
	if getContext(c).User.ID != 0 {
		resp403(c)
		return
	}
	if c.Query("stopsign") != "1" {
		u, _ := tryBotnets(c)
		if u != "" {
			simple(c, getSimpleByFilename("register/elmo.html"), nil, map[string]interface{}{
				"Username": u,
			})
			return
		}
	}
	registerResp(c)
}

func registerSubmit(c *gin.Context) {
	if getContext(c).User.ID != 0 {
		resp403(c)
		return
	}
	// check registrations are enabled
	if !registrationsEnabled() {
		registerResp(c, errorMessage{T(c, "Sorry, it's not possible to register at the moment. Please try again later.")})
		return
	}

	// check username is valid by our criteria
	username := strings.TrimSpace(c.PostForm("username"))
	if !usernameRegex.MatchString(username) {
		registerResp(c, errorMessage{T(c, "Your username must contain alphanumerical characters, spaces, or any of _[]-")})
		return
	}

	// check whether an username is e.g. cookiezi, shigetora, peppy, wubwoofwolf, loctav
	if in(strings.ToLower(username), forbiddenUsernames) {
		registerResp(c, errorMessage{T(c, "You're not allowed to register with that username.")})
		return
	}

	// check email is valid
	if !govalidator.IsEmail(c.PostForm("email")) {
		registerResp(c, errorMessage{T(c, "Please pass a valid email address.")})
		return
	}

	// passwords check (too short/too common)
	if x := validatePassword(c.PostForm("password")); x != "" {
		registerResp(c, errorMessage{T(c, x)})
		return
	}

	// usernames with both _ and spaces are not allowed
	if strings.Contains(username, "_") && strings.Contains(username, " ") {
		registerResp(c, errorMessage{T(c, "An username can't contain both underscores and spaces.")})
		return
	}

	// check whether username already exists
	if db.QueryRow("SELECT 1 FROM users WHERE username_safe = ?", safeUsername(username)).
		Scan(new(int)) != sql.ErrNoRows {
		registerResp(c, errorMessage{T(c, "An user with that username already exists!")})
		return
	}

	// check whether an user with that email already exists
	if db.QueryRow("SELECT 1 FROM users WHERE email = ?", c.PostForm("email")).
		Scan(new(int)) != sql.ErrNoRows {
		registerResp(c, errorMessage{T(c, "An user with that email address already exists!")})
		return
	}

	// recaptcha verify
	if config.RecaptchaPrivate != "" && !recaptchaCheck(c) {
		registerResp(c, errorMessage{T(c, "Captcha is invalid.")})
		return
	}

	uMulti, criteria := tryBotnets(c)
	if criteria != "" {
		schiavo.CMs.Send(
			fmt.Sprintf(
				"User **%s** registered with the same %s as %s (%s/u/%s). **POSSIBLE MULTIACCOUNT!!!**. Waiting for ingame verification...",
				username, criteria, uMulti, config.BaseURL, url.QueryEscape(uMulti),
			),
		)
	}

	// The actual registration.
	pass, err := generatePassword(c.PostForm("password"))
	if err != nil {
		resp500(c)
		return
	}

	res, err := db.Exec(`INSERT INTO users(username, username_safe, password_md5, salt, email, register_datetime, privileges, password_version)
							  VALUES (?,        ?,             ?,            '',   ?,     ?,                 ?,          2);`,
		username, safeUsername(username), pass, c.PostForm("email"), time.Now().Unix(), common.UserPrivilegePendingVerification)
	if err != nil {
		registerResp(c, errorMessage{T(c, "Whoops, an error slipped in. You might have been registered, though. I don't know.")})
		return
	}
	lid, _ := res.LastInsertId()

	db.Exec("INSERT INTO `users_stats`(id, username, user_color, user_style, ranked_score_std, playcount_std, total_score_std, ranked_score_taiko, playcount_taiko, total_score_taiko, ranked_score_ctb, playcount_ctb, total_score_ctb, ranked_score_mania, playcount_mania, total_score_mania) VALUES (?, ?, 'black', '', 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);", lid, username)
	db.Exec("INSERT INTO `users_stats_relax` (id) VALUES (?)", lid)

	schiavo.CMs.Send(fmt.Sprintf("User (**%s** | %s) registered from %s", username, c.PostForm("email"), clientIP(c)))

	setYCookie(int(lid), c)
	logIP(c, int(lid))

	rd.Incr("ripple:registered_users")

	addMessage(c, successMessage{T(c, "You have been successfully registered on Ripple! You now need to verify your account.")})
	getSession(c).Save()
	c.Redirect(302, "/register/verify?u="+strconv.Itoa(int(lid)))
}

func registerResp(c *gin.Context, messages ...message) {
	resp(c, 200, "register/register.html", &baseTemplateData{
		TitleBar:  "Register",
		KyutGrill: "register.jpg",
		Scripts:   []string{"https://www.google.com/recaptcha/api.js"},
		Messages:  messages,
		FormData:  normaliseURLValues(c.Request.PostForm),
	})
}

func registrationsEnabled() bool {
	var enabled bool
	db.QueryRow("SELECT value_int FROM system_settings WHERE name = 'registrations_enabled'").Scan(&enabled)
	return enabled
}

func verifyAccount(c *gin.Context) {
	if getContext(c).User.ID != 0 {
		resp403(c)
		return
	}

	i, ret := checkUInQS(c)
	if ret {
		return
	}

	sess := getSession(c)
	var rPrivileges uint64
	db.Get(&rPrivileges, "SELECT privileges FROM users WHERE id = ?", i)
	if common.UserPrivileges(rPrivileges)&common.UserPrivilegePendingVerification == 0 {
		addMessage(c, warningMessage{T(c, "Nope.")})
		sess.Save()
		c.Redirect(302, "/")
		return
	}

	resp(c, 200, "register/verify.html", &baseTemplateData{
		TitleBar:       "Verify account",
		HeadingOnRight: true,
		KyutGrill:      "welcome.jpg",
	})
}

func welcome(c *gin.Context) {
	if getContext(c).User.ID != 0 {
		resp403(c)
		return
	}

	i, ret := checkUInQS(c)
	if ret {
		return
	}

	var rPrivileges uint64
	db.Get(&rPrivileges, "SELECT privileges FROM users WHERE id = ?", i)
	if common.UserPrivileges(rPrivileges)&common.UserPrivilegePendingVerification > 0 {
		c.Redirect(302, "/register/verify?u="+c.Query("u"))
		return
	}

	t := T(c, "Welcome!")
	if common.UserPrivileges(rPrivileges)&common.UserPrivilegeNormal == 0 {
		// if the user has no UserNormal, it means they're banned = they multiaccounted
		t = T(c, "Welcome back!")
	}

	resp(c, 200, "register/welcome.html", &baseTemplateData{
		TitleBar:       t,
		HeadingOnRight: true,
		KyutGrill:      "welcome.jpg",
	})
}

// Check User In Query Is Same As User In Y Cookie
func checkUInQS(c *gin.Context) (int, bool) {
	sess := getSession(c)

	i, _ := strconv.Atoi(c.Query("u"))
	y, _ := c.Cookie("y")
	err := db.QueryRow("SELECT 1 FROM identity_tokens WHERE token = ? AND userid = ?", y, i).Scan(new(int))
	if err == sql.ErrNoRows {
		addMessage(c, warningMessage{T(c, "Nope.")})
		sess.Save()
		c.Redirect(302, "/")
		return 0, true
	}
	return i, false
}

func tryBotnets(c *gin.Context) (string, string) {
	var username string

	err := db.QueryRow("SELECT u.username FROM ip_user i INNER JOIN users u ON u.id = i.userid WHERE i.ip = ?", clientIP(c)).Scan(&username)
	if err != nil {
		if err != sql.ErrNoRows {
			c.Error(err)
		}
		return "", ""
	}
	if username != "" {
		return username, "IP"
	}

	cook, _ := c.Cookie("y")
	err = db.QueryRow("SELECT u.username FROM identity_tokens i INNER JOIN users u ON u.id = i.userid WHERE i.token = ?",
		cook).Scan(&username)
	if err != nil {
		if err != sql.ErrNoRows {
			c.Error(err)
		}
		return "", ""
	}
	if username != "" {
		return username, "username"
	}

	return "", ""
}

func in(s string, ss []string) bool {
	for _, x := range ss {
		if x == s {
			return true
		}
	}
	return false
}

var usernameRegex = regexp.MustCompile(`^[A-Za-z0-9 _\[\]-]{2,15}$`)
var forbiddenUsernames = []string{
	"peppy",
	"rrtyui",
	"cookiezi",
	"azer",
	"loctav",
	"banchobot",
	"happystick",
	"doomsday",
	"sharingan33",
	"andrea",
	"cptnxn",
	"reimu-desu",
	"hvick225",
	"_index",
	"my aim sucks",
	"kynan",
	"rafis",
	"sayonara-bye",
	"thelewa",
	"wubwoofwolf",
	"millhioref",
	"tom94",
	"tillerino",
	"clsw",
	"spectator",
	"exgon",
	"axarious",
	"angelsim",
	"recia",
	"nara",
	"emperorpenguin83",
	"bikko",
	"xilver",
	"vettel",
	"kuu01",
	"_yu68",
	"tasuke912",
	"dusk",
	"ttobas",
	"velperk",
	"jakads",
	"jhlee0133",
	"abcdullah",
	"yuko-",
	"entozer",
	"hdhr",
	"ekoro",
	"snowwhite",
	"osuplayer111",
	"musty",
	"nero",
	"elysion",
	"ztrot",
	"koreapenguin",
	"fort",
	"asphyxia",
	"niko",
	"shigetora",
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/bmizerany/mc"
	"github.com/elazarl/goproxy"
	"github.com/jackc/pgx"
	"github.com/pquerna/ffjson/ffjson"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strings"
)

func checkError(err error) {
	if err != nil {
		log.Fatalf("Error: %s", err)
	}
}

func UrlHasPrefix(prefix string) goproxy.ReqConditionFunc {
	return func(req *http.Request, ctx *goproxy.ProxyCtx) bool {
		return req.Method == http.MethodGet && strings.HasPrefix(req.URL.Path, prefix)
	}
}

func PathIs(path string) goproxy.ReqConditionFunc {
	return func(req *http.Request, ctx *goproxy.ProxyCtx) bool {
		return req.Method == http.MethodGet && req.URL.Path == path
	}
}

var pool *pgx.ConnPool

func afterConnect(conn *pgx.Conn) (err error) {
	_, err = conn.Prepare("getPackage", `
		SELECT name, url FROM packages WHERE name = $1
  `)
	if err != nil {
		return
	}

	return
}

func main() {
	var err error

	memcached_url := os.Getenv("MEMCACHEDCLOUD_SERVERS")
	memcached_username := os.Getenv("MEMCACHEDCLOUD_USERNAME")
	memcached_password := os.Getenv("MEMCACHEDCLOUD_PASSWORD")

	if memcached_url == "" {
		memcached_url = "localhost:11211"
	}

	cn, err := mc.Dial("tcp", memcached_url)
	if err != nil {
		log.Fatalf("Memcached connection error: %s", err)
	}

	if memcached_url != "" && memcached_password != "" {
		err = cn.Auth(memcached_username, memcached_password)
		if err != nil {
			log.Fatalf("Memcached auth error: %s", err)
		}
	}

	pgxcfg, err := pgx.ParseURI(os.Getenv("DATABASE_URL"))
	if err != nil {
		log.Fatalf("Parse URI error: %s", err)
	}
	pool, err = pgx.NewConnPool(pgx.ConnPoolConfig{
		ConnConfig:     pgxcfg,
		MaxConnections: 20,
		AfterConnect:   afterConnect,
	})
	if err != nil {
		log.Fatalf("Connection error: %s", err)
	}
	defer pool.Close()

	port := os.Getenv("PORT")

	if port == "" {
		port = "3000"
	}

	binary, err := exec.LookPath("node")
	checkError(err)
	cmd := exec.Command(binary, "--expose_gc", "index.js")
	env := os.Environ()
	env = append([]string{"PORT=3001"}, env...)
	cmd.Env = env
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err = cmd.Start()
	checkError(err)
	defer func() {
		err = cmd.Wait()
		checkError(err)
	}()

	proxy := goproxy.NewProxyHttpServer()

	proxy.Verbose = false

	proxy.NonproxyHandler = http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		req.Host = "bower.herokuapp.com"
		req.URL.Scheme = "http"
		req.URL.Host = "localhost:3001"
		proxy.ServeHTTP(w, req)
	})

	proxy.OnRequest(PathIs("/packages")).DoFunc(
		func(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
			var val string

			val, _, _, err = cn.Get("packages")

			switch err {
			case nil:
				fmt.Println("MEMCACHED FROM GO SERVER")
				return r, goproxy.NewResponse(r, "application/json", http.StatusOK, val)
			default:
				return r, nil
			}

			return r, nil
		})

	proxy.OnRequest(UrlHasPrefix("/packages/")).DoFunc(
		func(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
			var name string
			var url string
			elements := strings.Split(r.URL.Path, "/")
			packageName := elements[len(elements)-1]
			err := pool.QueryRow("getPackage", packageName).Scan(&name, &url)
			switch err {
			case nil:
				result := map[string]string{"name": name, "url": url}
				resultByteArray, _ := ffjson.Marshal(result)
				return r, goproxy.NewResponse(r, "application/json", http.StatusOK, string(resultByteArray))
			case pgx.ErrNoRows:
				return r, goproxy.NewResponse(r, "text/html", http.StatusNotFound, "Package not found")
			default:
				return r, goproxy.NewResponse(r, "text/html", http.StatusInternalServerError, "Internal server error")
			}
		})

	fmt.Println("Starting web server at port", port)

	log.Fatal(http.ListenAndServe(":"+port, proxy))
}
<commit_msg>Fix search endpoint<commit_after>package main

import (
	"fmt"
	"github.com/bmizerany/mc"
	"github.com/elazarl/goproxy"
	"github.com/jackc/pgx"
	"github.com/pquerna/ffjson/ffjson"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strings"
)

func checkError(err error) {
	if err != nil {
		log.Fatalf("Error: %s", err)
	}
}

func UrlHasPrefix(prefix string) goproxy.ReqConditionFunc {
	return func(req *http.Request, ctx *goproxy.ProxyCtx) bool {
		return req.Method == http.MethodGet && strings.HasPrefix(req.URL.Path, prefix) && !strings.HasPrefix(req.URL.Path, "/packages/search/")
	}
}

func PathIs(path string) goproxy.ReqConditionFunc {
	return func(req *http.Request, ctx *goproxy.ProxyCtx) bool {
		return req.Method == http.MethodGet && req.URL.Path == path
	}
}

var pool *pgx.ConnPool

func afterConnect(conn *pgx.Conn) (err error) {
	_, err = conn.Prepare("getPackage", `
		SELECT name, url FROM packages WHERE name = $1
  `)
	if err != nil {
		return
	}

	return
}

func main() {
	var err error

	memcached_url := os.Getenv("MEMCACHEDCLOUD_SERVERS")
	memcached_username := os.Getenv("MEMCACHEDCLOUD_USERNAME")
	memcached_password := os.Getenv("MEMCACHEDCLOUD_PASSWORD")

	if memcached_url == "" {
		memcached_url = "localhost:11211"
	}

	cn, err := mc.Dial("tcp", memcached_url)
	if err != nil {
		log.Fatalf("Memcached connection error: %s", err)
	}

	if memcached_url != "" && memcached_password != "" {
		err = cn.Auth(memcached_username, memcached_password)
		if err != nil {
			log.Fatalf("Memcached auth error: %s", err)
		}
	}

	pgxcfg, err := pgx.ParseURI(os.Getenv("DATABASE_URL"))
	if err != nil {
		log.Fatalf("Parse URI error: %s", err)
	}
	pool, err = pgx.NewConnPool(pgx.ConnPoolConfig{
		ConnConfig:     pgxcfg,
		MaxConnections: 20,
		AfterConnect:   afterConnect,
	})
	if err != nil {
		log.Fatalf("Connection error: %s", err)
	}
	defer pool.Close()

	port := os.Getenv("PORT")

	if port == "" {
		port = "3000"
	}

	binary, err := exec.LookPath("node")
	checkError(err)
	cmd := exec.Command(binary, "--expose_gc", "index.js")
	env := os.Environ()
	env = append([]string{"PORT=3001"}, env...)
	cmd.Env = env
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err = cmd.Start()
	checkError(err)
	defer func() {
		err = cmd.Wait()
		checkError(err)
	}()

	proxy := goproxy.NewProxyHttpServer()

	proxy.Verbose = false

	proxy.NonproxyHandler = http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		req.Host = "bower.herokuapp.com"
		req.URL.Scheme = "http"
		req.URL.Host = "localhost:3001"
		proxy.ServeHTTP(w, req)
	})

	proxy.OnRequest(PathIs("/packages")).DoFunc(
		func(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
			var val string

			val, _, _, err = cn.Get("packages")

			switch err {
			case nil:
				fmt.Println("MEMCACHED FROM GO SERVER")
				return r, goproxy.NewResponse(r, "application/json", http.StatusOK, val)
			default:
				return r, nil
			}

			return r, nil
		})

	proxy.OnRequest(UrlHasPrefix("/packages/")).DoFunc(
		func(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
			var name string
			var url string
			elements := strings.Split(r.URL.Path, "/")
			packageName := elements[len(elements)-1]
			err := pool.QueryRow("getPackage", packageName).Scan(&name, &url)
			switch err {
			case nil:
				result := map[string]string{"name": name, "url": url}
				resultByteArray, _ := ffjson.Marshal(result)
				return r, goproxy.NewResponse(r, "application/json", http.StatusOK, string(resultByteArray))
			case pgx.ErrNoRows:
				return r, goproxy.NewResponse(r, "text/html", http.StatusNotFound, "Package not found")
			default:
				return r, goproxy.NewResponse(r, "text/html", http.StatusInternalServerError, "Internal server error")
			}
		})

	fmt.Println("Starting web server at port", port)

	log.Fatal(http.ListenAndServe(":"+port, proxy))
}
<|endoftext|>


<commit_before>// Copyright 2011 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.

// Package crypto collects common cryptographic constants.
package crypto

import (
	"hash"
	"io"
	"strconv"
)

// Hash identifies a cryptographic hash function that is implemented in another
// package.
type Hash uint

// HashFunc simply returns the value of h so that Hash implements SignerOpts.
func (h Hash) HashFunc() Hash {
	return h
}

const (
	MD4       Hash = 1 + iota // import code.google.com/p/go.crypto/md4
	MD5                       // import crypto/md5
	SHA1                      // import crypto/sha1
	SHA224                    // import crypto/sha256
	SHA256                    // import crypto/sha256
	SHA384                    // import crypto/sha512
	SHA512                    // import crypto/sha512
	MD5SHA1                   // no implementation; MD5+SHA1 used for TLS RSA
	RIPEMD160                 // import code.google.com/p/go.crypto/ripemd160
	maxHash
)

var digestSizes = []uint8{
	MD4:       16,
	MD5:       16,
	SHA1:      20,
	SHA224:    28,
	SHA256:    32,
	SHA384:    48,
	SHA512:    64,
	MD5SHA1:   36,
	RIPEMD160: 20,
}

// Size returns the length, in bytes, of a digest resulting from the given hash
// function. It doesn't require that the hash function in question be linked
// into the program.
func (h Hash) Size() int {
	if h > 0 && h < maxHash {
		return int(digestSizes[h])
	}
	panic("crypto: Size of unknown hash function")
}

var hashes = make([]func() hash.Hash, maxHash)

// New returns a new hash.Hash calculating the given hash function. New panics
// if the hash function is not linked into the binary.
func (h Hash) New() hash.Hash {
	if h > 0 && h < maxHash {
		f := hashes[h]
		if f != nil {
			return f()
		}
	}
	panic("crypto: requested hash function #" + strconv.Itoa(int(h)) + " is unavailable")
}

// Available reports whether the given hash function is linked into the binary.
func (h Hash) Available() bool {
	return h < maxHash && hashes[h] != nil
}

// RegisterHash registers a function that returns a new instance of the given
// hash function. This is intended to be called from the init function in
// packages that implement hash functions.
func RegisterHash(h Hash, f func() hash.Hash) {
	if h >= maxHash {
		panic("crypto: RegisterHash of unknown hash function")
	}
	hashes[h] = f
}

// PublicKey represents a public key using an unspecified algorithm.
type PublicKey interface{}

// PrivateKey represents a private key using an unspecified algorithm.
type PrivateKey interface{}

// Signer is an interface for an opaque private key that can be used for
// signing operations. For example, an RSA key kept in a hardware module.
type Signer interface {
	// Public returns the public key corresponding to the opaque,
	// private key.
	Public() PublicKey

	// Sign signs msg with the private key, possibly using entropy from
	// rand. For an RSA key, the resulting signature should be either a
	// PKCS#1 v1.5 or PSS signature (as indicated by opts). For an (EC)DSA
	// key, it should be a DER-serialised, ASN.1 signature structure.
	//
	// Hash implements the SignerOpts interface and, in most cases, one can
	// simply pass in the hash function used as opts. Sign may also attempt
	// to type assert opts to other types in order to obtain algorithm
	// specific values. See the documentation in each package for details.
	Sign(rand io.Reader, msg []byte, opts SignerOpts) (signature []byte, err error)
}

// SignerOpts contains options for signing with a Signer.
type SignerOpts interface {
	// HashFunc returns an identifier for the hash function used to produce
	// the message passed to Signer.Sign, or else zero to indicate that no
	// hashing was done.
	HashFunc() Hash
}
<commit_msg>crypto: Add SHA3 functions in go.crypto/sha3 to the Hash enum.<commit_after>// Copyright 2011 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.

// Package crypto collects common cryptographic constants.
package crypto

import (
	"hash"
	"io"
	"strconv"
)

// Hash identifies a cryptographic hash function that is implemented in another
// package.
type Hash uint

// HashFunc simply returns the value of h so that Hash implements SignerOpts.
func (h Hash) HashFunc() Hash {
	return h
}

const (
	MD4       Hash = 1 + iota // import code.google.com/p/go.crypto/md4
	MD5                       // import crypto/md5
	SHA1                      // import crypto/sha1
	SHA224                    // import crypto/sha256
	SHA256                    // import crypto/sha256
	SHA384                    // import crypto/sha512
	SHA512                    // import crypto/sha512
	MD5SHA1                   // no implementation; MD5+SHA1 used for TLS RSA
	RIPEMD160                 // import code.google.com/p/go.crypto/ripemd160
	SHA3_224                  // import code.google.com/p/go.crypto/sha3
	SHA3_256                  // import code.google.com/p/go.crypto/sha3
	SHA3_384                  // import code.google.com/p/go.crypto/sha3
	SHA3_512                  // import code.google.com/p/go.crypto/sha3
	maxHash
)

var digestSizes = []uint8{
	MD4:       16,
	MD5:       16,
	SHA1:      20,
	SHA224:    28,
	SHA256:    32,
	SHA384:    48,
	SHA512:    64,
	SHA3_224:  28,
	SHA3_256:  32,
	SHA3_384:  48,
	SHA3_512:  64,
	MD5SHA1:   36,
	RIPEMD160: 20,
}

// Size returns the length, in bytes, of a digest resulting from the given hash
// function. It doesn't require that the hash function in question be linked
// into the program.
func (h Hash) Size() int {
	if h > 0 && h < maxHash {
		return int(digestSizes[h])
	}
	panic("crypto: Size of unknown hash function")
}

var hashes = make([]func() hash.Hash, maxHash)

// New returns a new hash.Hash calculating the given hash function. New panics
// if the hash function is not linked into the binary.
func (h Hash) New() hash.Hash {
	if h > 0 && h < maxHash {
		f := hashes[h]
		if f != nil {
			return f()
		}
	}
	panic("crypto: requested hash function #" + strconv.Itoa(int(h)) + " is unavailable")
}

// Available reports whether the given hash function is linked into the binary.
func (h Hash) Available() bool {
	return h < maxHash && hashes[h] != nil
}

// RegisterHash registers a function that returns a new instance of the given
// hash function. This is intended to be called from the init function in
// packages that implement hash functions.
func RegisterHash(h Hash, f func() hash.Hash) {
	if h >= maxHash {
		panic("crypto: RegisterHash of unknown hash function")
	}
	hashes[h] = f
}

// PublicKey represents a public key using an unspecified algorithm.
type PublicKey interface{}

// PrivateKey represents a private key using an unspecified algorithm.
type PrivateKey interface{}

// Signer is an interface for an opaque private key that can be used for
// signing operations. For example, an RSA key kept in a hardware module.
type Signer interface {
	// Public returns the public key corresponding to the opaque,
	// private key.
	Public() PublicKey

	// Sign signs msg with the private key, possibly using entropy from
	// rand. For an RSA key, the resulting signature should be either a
	// PKCS#1 v1.5 or PSS signature (as indicated by opts). For an (EC)DSA
	// key, it should be a DER-serialised, ASN.1 signature structure.
	//
	// Hash implements the SignerOpts interface and, in most cases, one can
	// simply pass in the hash function used as opts. Sign may also attempt
	// to type assert opts to other types in order to obtain algorithm
	// specific values. See the documentation in each package for details.
	Sign(rand io.Reader, msg []byte, opts SignerOpts) (signature []byte, err error)
}

// SignerOpts contains options for signing with a Signer.
type SignerOpts interface {
	// HashFunc returns an identifier for the hash function used to produce
	// the message passed to Signer.Sign, or else zero to indicate that no
	// hashing was done.
	HashFunc() Hash
}
<|endoftext|>