<commit_before>// ----------------------------------------------------------------------------
//
//     ***     AUTO GENERATED CODE    ***    AUTO GENERATED CODE     ***
//
// ----------------------------------------------------------------------------
//
//     This file is automatically generated by Magic Modules and manual
//     changes will be clobbered when the file is regenerated.
//
//     Please read more about how to change this file in
//     .github/CONTRIBUTING.md.
//
// ----------------------------------------------------------------------------

package google

import "reflect"

func GetComputeGlobalForwardingRuleCaiObject(d TerraformResourceData, config *Config) (Asset, error) {
	name, err := assetName(d, config, "//compute.googleapis.com/projects/{{project}}/global/forwardingRules/{{name}}")
	if err != nil {
		return Asset{}, err
	}
	if obj, err := GetComputeGlobalForwardingRuleApiObject(d, config); err == nil {
		return Asset{
			Name: name,
			Type: "compute.googleapis.com/GlobalForwardingRule",
			Resource: &AssetResource{
				Version:              "v1",
				DiscoveryDocumentURI: "https://www.googleapis.com/discovery/v1/apis/compute/v1/rest",
				DiscoveryName:        "GlobalForwardingRule",
				Data:                 obj,
			},
		}, nil
	} else {
		return Asset{}, err
	}
}

func GetComputeGlobalForwardingRuleApiObject(d TerraformResourceData, config *Config) (map[string]interface{}, error) {
	obj := make(map[string]interface{})
	descriptionProp, err := expandComputeGlobalForwardingRuleDescription(d.Get("description"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("description"); !isEmptyValue(reflect.ValueOf(descriptionProp)) && (ok || !reflect.DeepEqual(v, descriptionProp)) {
		obj["description"] = descriptionProp
	}
	IPAddressProp, err := expandComputeGlobalForwardingRuleIPAddress(d.Get("ip_address"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("ip_address"); !isEmptyValue(reflect.ValueOf(IPAddressProp)) && (ok || !reflect.DeepEqual(v, IPAddressProp)) {
		obj["IPAddress"] = IPAddressProp
	}
	IPProtocolProp, err := expandComputeGlobalForwardingRuleIPProtocol(d.Get("ip_protocol"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("ip_protocol"); !isEmptyValue(reflect.ValueOf(IPProtocolProp)) && (ok || !reflect.DeepEqual(v, IPProtocolProp)) {
		obj["IPProtocol"] = IPProtocolProp
	}
	ipVersionProp, err := expandComputeGlobalForwardingRuleIpVersion(d.Get("ip_version"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("ip_version"); !isEmptyValue(reflect.ValueOf(ipVersionProp)) && (ok || !reflect.DeepEqual(v, ipVersionProp)) {
		obj["ipVersion"] = ipVersionProp
	}
	nameProp, err := expandComputeGlobalForwardingRuleName(d.Get("name"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("name"); !isEmptyValue(reflect.ValueOf(nameProp)) && (ok || !reflect.DeepEqual(v, nameProp)) {
		obj["name"] = nameProp
	}
	portRangeProp, err := expandComputeGlobalForwardingRulePortRange(d.Get("port_range"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("port_range"); !isEmptyValue(reflect.ValueOf(portRangeProp)) && (ok || !reflect.DeepEqual(v, portRangeProp)) {
		obj["portRange"] = portRangeProp
	}
	targetProp, err := expandComputeGlobalForwardingRuleTarget(d.Get("target"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("target"); !isEmptyValue(reflect.ValueOf(targetProp)) && (ok || !reflect.DeepEqual(v, targetProp)) {
		obj["target"] = targetProp
	}

	return obj, nil
}

func expandComputeGlobalForwardingRuleDescription(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeGlobalForwardingRuleIPAddress(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeGlobalForwardingRuleIPProtocol(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeGlobalForwardingRuleIpVersion(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeGlobalForwardingRuleName(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeGlobalForwardingRulePortRange(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeGlobalForwardingRuleTarget(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}
<commit_msg>Add support for INTERNAL_SELF_MANAGED backend service<commit_after>// ----------------------------------------------------------------------------
//
//     ***     AUTO GENERATED CODE    ***    AUTO GENERATED CODE     ***
//
// ----------------------------------------------------------------------------
//
//     This file is automatically generated by Magic Modules and manual
//     changes will be clobbered when the file is regenerated.
//
//     Please read more about how to change this file in
//     .github/CONTRIBUTING.md.
//
// ----------------------------------------------------------------------------

package google

import "reflect"

func GetComputeGlobalForwardingRuleCaiObject(d TerraformResourceData, config *Config) (Asset, error) {
	name, err := assetName(d, config, "//compute.googleapis.com/projects/{{project}}/global/forwardingRules/{{name}}")
	if err != nil {
		return Asset{}, err
	}
	if obj, err := GetComputeGlobalForwardingRuleApiObject(d, config); err == nil {
		return Asset{
			Name: name,
			Type: "compute.googleapis.com/GlobalForwardingRule",
			Resource: &AssetResource{
				Version:              "v1",
				DiscoveryDocumentURI: "https://www.googleapis.com/discovery/v1/apis/compute/v1/rest",
				DiscoveryName:        "GlobalForwardingRule",
				Data:                 obj,
			},
		}, nil
	} else {
		return Asset{}, err
	}
}

func GetComputeGlobalForwardingRuleApiObject(d TerraformResourceData, config *Config) (map[string]interface{}, error) {
	obj := make(map[string]interface{})
	descriptionProp, err := expandComputeGlobalForwardingRuleDescription(d.Get("description"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("description"); !isEmptyValue(reflect.ValueOf(descriptionProp)) && (ok || !reflect.DeepEqual(v, descriptionProp)) {
		obj["description"] = descriptionProp
	}
	IPAddressProp, err := expandComputeGlobalForwardingRuleIPAddress(d.Get("ip_address"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("ip_address"); !isEmptyValue(reflect.ValueOf(IPAddressProp)) && (ok || !reflect.DeepEqual(v, IPAddressProp)) {
		obj["IPAddress"] = IPAddressProp
	}
	IPProtocolProp, err := expandComputeGlobalForwardingRuleIPProtocol(d.Get("ip_protocol"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("ip_protocol"); !isEmptyValue(reflect.ValueOf(IPProtocolProp)) && (ok || !reflect.DeepEqual(v, IPProtocolProp)) {
		obj["IPProtocol"] = IPProtocolProp
	}
	ipVersionProp, err := expandComputeGlobalForwardingRuleIpVersion(d.Get("ip_version"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("ip_version"); !isEmptyValue(reflect.ValueOf(ipVersionProp)) && (ok || !reflect.DeepEqual(v, ipVersionProp)) {
		obj["ipVersion"] = ipVersionProp
	}
	loadBalancingSchemeProp, err := expandComputeGlobalForwardingRuleLoadBalancingScheme(d.Get("load_balancing_scheme"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("load_balancing_scheme"); !isEmptyValue(reflect.ValueOf(loadBalancingSchemeProp)) && (ok || !reflect.DeepEqual(v, loadBalancingSchemeProp)) {
		obj["loadBalancingScheme"] = loadBalancingSchemeProp
	}
	nameProp, err := expandComputeGlobalForwardingRuleName(d.Get("name"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("name"); !isEmptyValue(reflect.ValueOf(nameProp)) && (ok || !reflect.DeepEqual(v, nameProp)) {
		obj["name"] = nameProp
	}
	portRangeProp, err := expandComputeGlobalForwardingRulePortRange(d.Get("port_range"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("port_range"); !isEmptyValue(reflect.ValueOf(portRangeProp)) && (ok || !reflect.DeepEqual(v, portRangeProp)) {
		obj["portRange"] = portRangeProp
	}
	targetProp, err := expandComputeGlobalForwardingRuleTarget(d.Get("target"), d, config)
	if err != nil {
		return nil, err
	} else if v, ok := d.GetOkExists("target"); !isEmptyValue(reflect.ValueOf(targetProp)) && (ok || !reflect.DeepEqual(v, targetProp)) {
		obj["target"] = targetProp
	}

	return obj, nil
}

func expandComputeGlobalForwardingRuleDescription(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeGlobalForwardingRuleIPAddress(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeGlobalForwardingRuleIPProtocol(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeGlobalForwardingRuleIpVersion(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeGlobalForwardingRuleLoadBalancingScheme(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeGlobalForwardingRuleName(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeGlobalForwardingRulePortRange(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}

func expandComputeGlobalForwardingRuleTarget(v interface{}, d TerraformResourceData, config *Config) (interface{}, error) {
	return v, nil
}
<|endoftext|>


<commit_before>package metadata_manager

import (
	log "github.com/Sirupsen/logrus"
	"github.com/samuel/go-zookeeper/zk"
	// "github.com/golang/protobuf/proto"
	"fmt"
	"strings"
	"sync"
	"time"
)

type ZkNode struct {
	mgr  *MetadataManager
	stat *zk.Stat
	data []byte
	ns   Namespace
}

func (node *ZkNode) Delete() {
	node.mgr.zkConn.Delete(node.ns.GetZKPath(), node.stat.Version)

}
func (node *ZkNode) String() string {
	return fmt.Sprintf("<%s> -> %v", node.ns.GetZKPath(), node.data)
}
func (node *ZkNode) GetData() []byte {
	return node.data
}
func (node *ZkNode) GetLock() *zk.Lock {
	zkLock := zk.NewLock(node.mgr.zkConn, node.ns.GetZKPath(), zk.WorldACL(zk.PermAll))
	return zkLock
}
func (node *ZkNode) SetData(data []byte) {
	var err error
	log.Info("Persisting data")
	if node.stat != nil {
		node.stat, err = node.mgr.zkConn.Set(node.ns.GetZKPath(), data, node.stat.Version)
		if err != nil {
			log.Panic("Error persisting data: ", err)
		}
	} else {
		_, err = node.mgr.zkConn.Create(node.ns.GetZKPath(), data, 0, zk.WorldACL(zk.PermAll))
		if err != nil {
			log.Panic("Error persisting data: ", err)
		}
		node.data, node.stat, err = node.mgr.zkConn.Get(node.ns.GetZKPath())
		if err != nil {
			log.Panic("Error persisting data: ", err)
		}
	}
}
func (node *ZkNode) GetChildren() []*ZkNode {
	return node.mgr.getChildren(node.ns)
}

func (node *ZkNode) GetChildrenW() ([]*ZkNode, <-chan zk.Event) {
	return node.mgr.getChildrenW(node.ns)
}

func (node *ZkNode) MakeEmptyChild(name string) *ZkNode {
	if strings.Contains(name, "/") {
		panic("Error, name of subnode cannot contain /")
	}
	ns := makeSubSpace(node.ns, name)
	newNode := &ZkNode{
		mgr: node.mgr,
		ns:  ns,
	}
	return newNode
}
func (node *ZkNode) MakeChild(name string, ephemeral bool) (*ZkNode, error) {
	if strings.Contains(name, "/") {
		panic("Error, name of subnode cannot contain /")
	}
	ns := makeSubSpace(node.ns, name)
	return node.mgr.makeNode(ns, ephemeral)
}

func (node *ZkNode) MakeChildWithData(name string, data []byte, ephemeral bool) (*ZkNode, error) {
	if strings.Contains(name, "/") {
		panic("Error, name of subnode cannot contain /")
	}
	ns := makeSubSpace(node.ns, name)
	return node.mgr.makeNodeWithData(ns, data, ephemeral)
}

func (node *ZkNode) GetChild(name string) (*ZkNode, error) {
	if strings.Contains(name, "/") {
		panic("Error, name of subnode cannot contain /")
	}
	ns := makeSubSpace(node.ns, name)
	return node.mgr.getNode(ns)
}

func (node *ZkNode) CreateChildIfNotExists(name string) {
	if strings.Contains(name, "/") {
		panic("Error, name of subnode cannot contain /")
	}
	ns := makeSubSpace(node.ns, name)
	node.mgr.CreateNSIfNotExists(ns, false)
}

type Namespace interface {
	GetComponents() []string
	GetZKPath() string
}
type baseNamespace struct {
}

// Base namespace should only ever return "" -- at least for Zookeeper
func (baseNamespace) GetComponents() []string {
	return []string{""}
}

// Base namespace should only ever return "" -- at least for Zookeeper
func (baseNamespace) GetZKPath() string {
	return "/"
}

type SubNamespace struct {
	parent    Namespace
	component string
}

// Components are read-only, so not pointer-receiver
func (ns SubNamespace) GetComponents() []string {
	return append(ns.parent.GetComponents(), ns.component)
}
func (ns SubNamespace) GetZKPath() string {
	return strings.Join(ns.GetComponents(), "/")
}
func makeSubSpace(ns Namespace, subSpaceName string) Namespace {
	return SubNamespace{parent: ns, component: subSpaceName}
}

type MetadataManager struct {
	framework   MetadataManagerFramework
	frameworkID string
	zkConn      *zk.Conn
	namespace   Namespace
	lock        *sync.Mutex
	zkLock      zk.Lock
}

func (mgr *MetadataManager) setup() {
	mgr.CreateNSIfNotExists(mgr.namespace, false)
}

func NewMetadataManager(frameworkID string, zookeepers []string) *MetadataManager {
	conn, _, err := zk.Connect(zookeepers, time.Second)
	if err != nil {
		log.Panic(err)
	}
	bns := baseNamespace{}
	ns := makeSubSpace(makeSubSpace(makeSubSpace(bns, "riak"), "frameworks"), frameworkID)
	lockPath := makeSubSpace(ns, "lock")
	zkLock := zk.NewLock(conn, lockPath.GetZKPath(), zk.WorldACL(zk.PermAll))

	manager := &MetadataManager{
		lock:        &sync.Mutex{},
		frameworkID: frameworkID,
		zkConn:      conn,
		namespace:   ns,
		zkLock:      *zkLock,
	}

	manager.setup()
	return manager
}
func (mgr *MetadataManager) createPathIfNotExists(path string, ephemeral bool) {
	splitString := strings.Split(path, "/")
	for idx := range splitString {
		if idx == 0 {
			continue
		}
		mgr.createIfNotExists(strings.Join(splitString[0:idx+1], "/"), ephemeral)
	}
}

func (mgr *MetadataManager) CreateNSIfNotExists(ns Namespace, ephemeral bool) {
	components := ns.GetComponents()
	for idx := range components {
		if idx == 0 {
			continue
		}
		mgr.createIfNotExists(strings.Join(components[0:idx+1], "/"), ephemeral)
	}
}
func (mgr *MetadataManager) createIfNotExists(path string, ephemeral bool) {
	exists, _, err := mgr.zkConn.Exists(path)
	if err != nil {
		log.Panic(err)
	}
	if !exists {
		var err error
		if ephemeral {
			_, err = mgr.zkConn.Create(path, nil, zk.FlagEphemeral, zk.WorldACL(zk.PermAll))
		} else {
			_, err = mgr.zkConn.Create(path, nil, 0, zk.WorldACL(zk.PermAll))
		}
		if err != nil {
			log.Panic(err)
		}
	}
}

// This subspaces the node in the "current working namespace"
func (mgr *MetadataManager) GetRootNode() *ZkNode {
	node, err := mgr.getNode(mgr.namespace)
	if err != nil {
		log.Panic("Could not get Root node")
	}
	return node
}

func (mgr *MetadataManager) getChildrenW(ns Namespace) ([]*ZkNode, <-chan zk.Event) {
	children, _, watchChan, err := mgr.zkConn.ChildrenW(ns.GetZKPath())
	if err != nil {
		log.Panic(err)
	}
	result := make([]*ZkNode, len(children))
	for idx, name := range children {
		result[idx], err = mgr.getNode(makeSubSpace(ns, name))
		if err != nil {
			log.Panic(err)
		}
	}
	return result, watchChan
}
func (mgr *MetadataManager) getChildren(ns Namespace) []*ZkNode {
	children, _, err := mgr.zkConn.Children(ns.GetZKPath())
	if err != nil {
		log.Panic(err)
	}
	result := make([]*ZkNode, len(children))
	for idx, name := range children {
		result[idx], err = mgr.getNode(makeSubSpace(ns, name))
		if err != nil {
			log.Panic(err)
		}
	}
	return result
}

func (mgr *MetadataManager) getNode(ns Namespace) (*ZkNode, error) {
	// Namespaces are also nodes
	data, stat, err := mgr.zkConn.Get(ns.GetZKPath())
	if err != nil {
		return nil, err
	}
	node := &ZkNode{
		mgr:  mgr,
		data: data,
		stat: stat,
		ns:   ns,
	}
	return node, nil
}

func (mgr *MetadataManager) makeNode(ns Namespace, ephemeral bool) (*ZkNode, error) {
	var flags int32
	if ephemeral {
		flags = zk.FlagEphemeral
	} else {
		flags = 0
	}
	// Namespaces are also nodes
	log.Info("Making node")
	_, err := mgr.zkConn.Create(ns.GetZKPath(), nil, flags, zk.WorldACL(zk.PermAll))
	if err != nil {
		log.Panic(err)
	}
	return mgr.getNode(ns)
}

func (mgr *MetadataManager) makeNodeWithData(ns Namespace, data []byte, ephemeral bool) (*ZkNode, error) {
	var flags int32
	if ephemeral {
		flags = zk.FlagEphemeral
	} else {
		flags = 0
	}
	// Namespaces are also nodes
	log.Info("Making node")
	_, err := mgr.zkConn.Create(ns.GetZKPath(), data, flags, zk.WorldACL(zk.PermAll))
	if err != nil {
		return nil, err
	}
	return mgr.getNode(ns)
}
<commit_msg>Start metadata manager for testing<commit_after>package metadata_manager

import (
	log "github.com/Sirupsen/logrus"
	"github.com/samuel/go-zookeeper/zk"
	// "github.com/golang/protobuf/proto"
	"fmt"
	"strings"
	"sync"
	"time"
)

// TODO: Convert ZKNode functions to all work around MetadataNode interface for better testing
type MetadataNode interface {
}
type ZkNode struct {
	mgr  *MetadataManager
	stat *zk.Stat
	data []byte
	ns   Namespace
}

func (node *ZkNode) Delete() {
	node.mgr.zkConn.Delete(node.ns.GetZKPath(), node.stat.Version)

}
func (node *ZkNode) String() string {
	return fmt.Sprintf("<%s> -> %v", node.ns.GetZKPath(), node.data)
}
func (node *ZkNode) GetData() []byte {
	return node.data
}
func (node *ZkNode) GetLock() *zk.Lock {
	zkLock := zk.NewLock(node.mgr.zkConn, node.ns.GetZKPath(), zk.WorldACL(zk.PermAll))
	return zkLock
}
func (node *ZkNode) SetData(data []byte) {
	var err error
	log.Info("Persisting data")
	if node.stat != nil {
		node.stat, err = node.mgr.zkConn.Set(node.ns.GetZKPath(), data, node.stat.Version)
		if err != nil {
			log.Panic("Error persisting data: ", err)
		}
	} else {
		_, err = node.mgr.zkConn.Create(node.ns.GetZKPath(), data, 0, zk.WorldACL(zk.PermAll))
		if err != nil {
			log.Panic("Error persisting data: ", err)
		}
		node.data, node.stat, err = node.mgr.zkConn.Get(node.ns.GetZKPath())
		if err != nil {
			log.Panic("Error persisting data: ", err)
		}
	}
}
func (node *ZkNode) GetChildren() []*ZkNode {
	return node.mgr.getChildren(node.ns)
}

func (node *ZkNode) GetChildrenW() ([]*ZkNode, <-chan zk.Event) {
	return node.mgr.getChildrenW(node.ns)
}

func (node *ZkNode) MakeEmptyChild(name string) *ZkNode {
	if strings.Contains(name, "/") {
		panic("Error, name of subnode cannot contain /")
	}
	ns := makeSubSpace(node.ns, name)
	newNode := &ZkNode{
		mgr: node.mgr,
		ns:  ns,
	}
	return newNode
}
func (node *ZkNode) MakeChild(name string, ephemeral bool) (*ZkNode, error) {
	if strings.Contains(name, "/") {
		panic("Error, name of subnode cannot contain /")
	}
	ns := makeSubSpace(node.ns, name)
	return node.mgr.makeNode(ns, ephemeral)
}

func (node *ZkNode) MakeChildWithData(name string, data []byte, ephemeral bool) (*ZkNode, error) {
	if strings.Contains(name, "/") {
		panic("Error, name of subnode cannot contain /")
	}
	ns := makeSubSpace(node.ns, name)
	return node.mgr.makeNodeWithData(ns, data, ephemeral)
}

func (node *ZkNode) GetChild(name string) (*ZkNode, error) {
	if strings.Contains(name, "/") {
		panic("Error, name of subnode cannot contain /")
	}
	ns := makeSubSpace(node.ns, name)
	return node.mgr.getNode(ns)
}

func (node *ZkNode) CreateChildIfNotExists(name string) {
	if strings.Contains(name, "/") {
		panic("Error, name of subnode cannot contain /")
	}
	ns := makeSubSpace(node.ns, name)
	node.mgr.CreateNSIfNotExists(ns, false)
}

type Namespace interface {
	GetComponents() []string
	GetZKPath() string
}
type baseNamespace struct {
}

// Base namespace should only ever return "" -- at least for Zookeeper
func (baseNamespace) GetComponents() []string {
	return []string{""}
}

// Base namespace should only ever return "" -- at least for Zookeeper
func (baseNamespace) GetZKPath() string {
	return "/"
}

type SubNamespace struct {
	parent    Namespace
	component string
}

// Components are read-only, so not pointer-receiver
func (ns SubNamespace) GetComponents() []string {
	return append(ns.parent.GetComponents(), ns.component)
}
func (ns SubNamespace) GetZKPath() string {
	return strings.Join(ns.GetComponents(), "/")
}
func makeSubSpace(ns Namespace, subSpaceName string) Namespace {
	return SubNamespace{parent: ns, component: subSpaceName}
}

type MetadataManager struct {
	framework   MetadataManagerFramework
	frameworkID string
	zkConn      *zk.Conn
	namespace   Namespace
	lock        *sync.Mutex
	zkLock      zk.Lock
}

func (mgr *MetadataManager) setup() {
	mgr.CreateNSIfNotExists(mgr.namespace, false)
}

func NewMetadataManager(frameworkID string, zookeepers []string) *MetadataManager {
	conn, _, err := zk.Connect(zookeepers, time.Second)
	if err != nil {
		log.Panic(err)
	}
	bns := baseNamespace{}
	ns := makeSubSpace(makeSubSpace(makeSubSpace(bns, "riak"), "frameworks"), frameworkID)
	lockPath := makeSubSpace(ns, "lock")
	zkLock := zk.NewLock(conn, lockPath.GetZKPath(), zk.WorldACL(zk.PermAll))

	manager := &MetadataManager{
		lock:        &sync.Mutex{},
		frameworkID: frameworkID,
		zkConn:      conn,
		namespace:   ns,
		zkLock:      *zkLock,
	}

	manager.setup()
	return manager
}
func (mgr *MetadataManager) createPathIfNotExists(path string, ephemeral bool) {
	splitString := strings.Split(path, "/")
	for idx := range splitString {
		if idx == 0 {
			continue
		}
		mgr.createIfNotExists(strings.Join(splitString[0:idx+1], "/"), ephemeral)
	}
}

func (mgr *MetadataManager) CreateNSIfNotExists(ns Namespace, ephemeral bool) {
	components := ns.GetComponents()
	for idx := range components {
		if idx == 0 {
			continue
		}
		mgr.createIfNotExists(strings.Join(components[0:idx+1], "/"), ephemeral)
	}
}
func (mgr *MetadataManager) createIfNotExists(path string, ephemeral bool) {
	exists, _, err := mgr.zkConn.Exists(path)
	if err != nil {
		log.Panic(err)
	}
	if !exists {
		var err error
		if ephemeral {
			_, err = mgr.zkConn.Create(path, nil, zk.FlagEphemeral, zk.WorldACL(zk.PermAll))
		} else {
			_, err = mgr.zkConn.Create(path, nil, 0, zk.WorldACL(zk.PermAll))
		}
		if err != nil {
			log.Panic(err)
		}
	}
}

// This subspaces the node in the "current working namespace"
func (mgr *MetadataManager) GetRootNode() *ZkNode {
	node, err := mgr.getNode(mgr.namespace)
	if err != nil {
		log.Panic("Could not get Root node")
	}
	return node
}

func (mgr *MetadataManager) getChildrenW(ns Namespace) ([]*ZkNode, <-chan zk.Event) {
	children, _, watchChan, err := mgr.zkConn.ChildrenW(ns.GetZKPath())
	if err != nil {
		log.Panic(err)
	}
	result := make([]*ZkNode, len(children))
	for idx, name := range children {
		result[idx], err = mgr.getNode(makeSubSpace(ns, name))
		if err != nil {
			log.Panic(err)
		}
	}
	return result, watchChan
}
func (mgr *MetadataManager) getChildren(ns Namespace) []*ZkNode {
	children, _, err := mgr.zkConn.Children(ns.GetZKPath())
	if err != nil {
		log.Panic(err)
	}
	result := make([]*ZkNode, len(children))
	for idx, name := range children {
		result[idx], err = mgr.getNode(makeSubSpace(ns, name))
		if err != nil {
			log.Panic(err)
		}
	}
	return result
}

func (mgr *MetadataManager) getNode(ns Namespace) (*ZkNode, error) {
	// Namespaces are also nodes
	data, stat, err := mgr.zkConn.Get(ns.GetZKPath())
	if err != nil {
		return nil, err
	}
	node := &ZkNode{
		mgr:  mgr,
		data: data,
		stat: stat,
		ns:   ns,
	}
	return node, nil
}

func (mgr *MetadataManager) makeNode(ns Namespace, ephemeral bool) (*ZkNode, error) {
	var flags int32
	if ephemeral {
		flags = zk.FlagEphemeral
	} else {
		flags = 0
	}
	// Namespaces are also nodes
	log.Info("Making node")
	_, err := mgr.zkConn.Create(ns.GetZKPath(), nil, flags, zk.WorldACL(zk.PermAll))
	if err != nil {
		log.Panic(err)
	}
	return mgr.getNode(ns)
}

func (mgr *MetadataManager) makeNodeWithData(ns Namespace, data []byte, ephemeral bool) (*ZkNode, error) {
	var flags int32
	if ephemeral {
		flags = zk.FlagEphemeral
	} else {
		flags = 0
	}
	// Namespaces are also nodes
	log.Info("Making node")
	_, err := mgr.zkConn.Create(ns.GetZKPath(), data, flags, zk.WorldACL(zk.PermAll))
	if err != nil {
		return nil, err
	}
	return mgr.getNode(ns)
}
<|endoftext|>


<commit_before>package attachments

import (
	"errors"
	"io"
	"path/filepath"
	"testing"
	"time"

	"github.com/keybase/client/go/chat/globals"
	"github.com/keybase/client/go/chat/storage"
	"github.com/keybase/client/go/chat/types"
	"github.com/keybase/client/go/kbtest"
	"github.com/keybase/client/go/libkb"
	"github.com/stretchr/testify/require"

	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
	"golang.org/x/net/context"
)

type mockStore struct {
	Store
	uploadFn func(context.Context, *UploadTask) (chat1.Asset, error)
}

func (m *mockStore) UploadAsset(ctx context.Context, task *UploadTask, encryptedOut io.Writer) (chat1.Asset, error) {
	return m.uploadFn(ctx, task)
}

type mockRemote struct {
	chat1.RemoteInterface
}

func (r mockRemote) GetS3Params(context.Context, chat1.ConversationID) (chat1.S3Params, error) {
	return chat1.S3Params{}, nil
}

func (r mockRemote) S3Sign(context.Context, chat1.S3SignArg) ([]byte, error) {
	return nil, nil
}

type mockActivityNotifier struct {
	types.ActivityNotifier
	startCh chan chat1.OutboxID
}

func newMockActivityNotifier() *mockActivityNotifier {
	return &mockActivityNotifier{
		startCh: make(chan chat1.OutboxID, 1000),
	}
}

func (a *mockActivityNotifier) AttachmentUploadStart(ctx context.Context, uid gregor1.UID,
	convID chat1.ConversationID, outboxID chat1.OutboxID) {
	a.startCh <- outboxID
}

func (a *mockActivityNotifier) AttachmentUploadProgress(ctx context.Context, uid gregor1.UID,
	convID chat1.ConversationID, outboxID chat1.OutboxID, bytesComplete, bytesTotal int64) {

}

type mockDeliverer struct {
	types.MessageDeliverer
	forceCh chan struct{}
}

func newMockDeliverer() *mockDeliverer {
	return &mockDeliverer{
		forceCh: make(chan struct{}, 1000),
	}
}

func (m *mockDeliverer) ForceDeliverLoop(context.Context) {
	m.forceCh <- struct{}{}
}

func (m *mockDeliverer) Stop(context.Context) chan struct{} {
	ch := make(chan struct{})
	close(ch)
	return ch
}

func TestAttachmentUploader(t *testing.T) {
	world := kbtest.NewChatMockWorld(t, "uploader", 1)
	defer world.Cleanup()

	u := world.GetUsers()[0]
	uid := gregor1.UID(u.User.GetUID().ToBytes())
	tc := world.Tcs[u.Username]
	g := globals.NewContext(tc.G, tc.ChatG)
	notifier := newMockActivityNotifier()
	store := &mockStore{}
	ri := mockRemote{}
	deliverer := newMockDeliverer()
	g.AttachmentURLSrv = types.DummyAttachmentHTTPSrv{}
	g.ActivityNotifier = notifier
	g.MessageDeliverer = deliverer
	getRi := func() chat1.RemoteInterface { return ri }
	cacheSize := 1
	uploader := NewUploader(g, store, NewS3Signer(getRi), getRi, cacheSize)
	convID := chat1.ConversationID([]byte{0, 1, 0})
	outboxID, err := storage.NewOutboxID()
	require.NoError(t, err)
	filename := "../testdata/ship.jpg"

	// Basic test to see if it works
	store.uploadFn = func(context.Context, *UploadTask) (chat1.Asset, error) {
		return chat1.Asset{}, nil
	}
	md, err := libkb.RandBytes(10)
	require.NoError(t, err)
	resChan, err := uploader.Register(context.TODO(), uid, convID, outboxID, "ship", filename, md, nil)
	require.NoError(t, err)
	uploadStartCheck := func(shouldHappen bool, outboxID chat1.OutboxID) {
		if shouldHappen {
			select {
			case obid := <-notifier.startCh:
				require.Equal(t, outboxID, obid)
			case <-time.After(20 * time.Second):
				require.Fail(t, "no start")
			}
		} else {
			select {
			case <-notifier.startCh:
				require.Fail(t, "start not supposed to happen")
			default:
			}
		}
	}
	deliverCheck := func(shouldHappen bool) {
		if shouldHappen {
			select {
			case <-deliverer.forceCh:
			case <-time.After(20 * time.Second):
				require.Fail(t, "no start")
			}
		} else {
			select {
			case <-deliverer.forceCh:
				require.Fail(t, "start not supposed to happen")
			default:
			}
		}
	}
	successCheck := func(cb types.AttachmentUploaderResultCb) {
		ch := cb.Wait()
		select {
		case res := <-ch:
			require.Nil(t, res.Error)
			require.Equal(t, md, res.Metadata)
			require.NotNil(t, res.Preview)
			require.Equal(t, "image/jpeg", res.Preview.MimeType)
			require.Equal(t, "image/jpeg", res.Object.MimeType)
		case <-time.After(20 * time.Second):
			require.Fail(t, "no upload")
		}
	}
	deliverCheck(true)
	uploadStartCheck(true, outboxID)
	successCheck(resChan)

	// Broken store
	outboxID, err = storage.NewOutboxID()
	require.NoError(t, err)
	store.uploadFn = func(context.Context, *UploadTask) (chat1.Asset, error) {
		return chat1.Asset{}, errors.New("i dont work")
	}
	resChan, err = uploader.Register(context.TODO(), uid, convID, outboxID, "ship", filename, md, nil)
	require.NoError(t, err)
	uploadStartCheck(true, outboxID)
	select {
	case res := <-resChan.Wait():
		require.NotNil(t, res.Error)
	case <-time.After(20 * time.Second):
		require.Fail(t, "no upload")
	}
	deliverCheck(true)

	// Retry after fixing store
	store.uploadFn = func(context.Context, *UploadTask) (chat1.Asset, error) {
		return chat1.Asset{}, nil
	}
	resChan, err = uploader.Retry(context.TODO(), outboxID)
	require.NoError(t, err)
	uploadStartCheck(true, outboxID)
	successCheck(resChan)
	deliverCheck(true)

	// Slow store to test concurrent retry
	outboxID, err = storage.NewOutboxID()
	require.NoError(t, err)
	slowCh := make(chan struct{})
	store.uploadFn = func(context.Context, *UploadTask) (chat1.Asset, error) {
		<-slowCh
		return chat1.Asset{}, nil
	}
	resChan, err = uploader.Register(context.TODO(), uid, convID, outboxID, "ship", filename, md, nil)
	require.NoError(t, err)
	uploadStartCheck(true, outboxID)
	deliverCheck(false)
	select {
	case <-resChan.Wait():
		require.Fail(t, "no res")
	default:
	}
	retryChan, err := uploader.Retry(context.TODO(), outboxID)
	require.NoError(t, err)
	uploadStartCheck(false, outboxID)
	close(slowCh)
	deliverCheck(true)
	// Should get results on both of these
	successCheck(retryChan)
	successCheck(resChan)

	uploader.Complete(context.TODO(), outboxID)
	_, _, err = uploader.Status(context.TODO(), outboxID)
	require.Error(t, err)

	// Test cancel
	outboxID, err = storage.NewOutboxID()
	require.NoError(t, err)
	slowCh = make(chan struct{})
	store.uploadFn = func(ctx context.Context, task *UploadTask) (chat1.Asset, error) {
		select {
		case <-slowCh:
		case <-ctx.Done():
			return chat1.Asset{}, ctx.Err()
		}
		return chat1.Asset{}, nil
	}
	resChan, err = uploader.Register(context.TODO(), uid, convID, outboxID, "ship", filename, md, nil)
	require.NoError(t, err)
	uploadStartCheck(true, outboxID)
	deliverCheck(false)
	select {
	case <-resChan.Wait():
		require.Fail(t, "no res")
	default:
	}
	require.NoError(t, uploader.Cancel(context.TODO(), outboxID))
	_, _, err = uploader.Status(context.TODO(), outboxID)
	require.Error(t, err)
	res := <-resChan.Wait()
	require.NotNil(t, res.Error)

	// verify uploadedPreviewsDir respects the cache size
	baseDir := uploader.getBaseDir()
	uploadedPreviews, err := filepath.Glob(filepath.Join(baseDir, uploadedPreviewsDir, "*"))
	require.NoError(t, err)
	require.Len(t, uploadedPreviews, 1)

	// verify uploadedFullsDir is respects the cache size
	uploadedFulls, err := filepath.Glob(filepath.Join(baseDir, uploadedFullsDir, "*"))
	require.NoError(t, err)
	require.Len(t, uploadedFulls, 1)
	mctx := kbtest.NewMetaContextForTest(*tc)

	// verify db nuke
	_, err = g.LocalDb.Nuke()
	require.NoError(t, err)
	err = uploader.OnDbNuke(mctx)
	require.NoError(t, err)

	uploadedPreviews, err = filepath.Glob(filepath.Join(baseDir, uploadedPreviewsDir, "*"))
	require.NoError(t, err)
	require.Zero(t, len(uploadedPreviews))

	uploadedFulls, err = filepath.Glob(filepath.Join(baseDir, uploadedFullsDir, "*"))
	require.NoError(t, err)
	require.Zero(t, len(uploadedFulls))
}
<commit_msg>fix uploader test flake (#23183)<commit_after>package attachments

import (
	"errors"
	"io"
	"path/filepath"
	"testing"
	"time"

	"github.com/keybase/client/go/chat/globals"
	"github.com/keybase/client/go/chat/storage"
	"github.com/keybase/client/go/chat/types"
	"github.com/keybase/client/go/kbtest"
	"github.com/keybase/client/go/libkb"
	"github.com/stretchr/testify/require"

	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
	"golang.org/x/net/context"
)

type mockStore struct {
	Store
	uploadFn func(context.Context, *UploadTask) (chat1.Asset, error)
}

func (m *mockStore) UploadAsset(ctx context.Context, task *UploadTask, encryptedOut io.Writer) (chat1.Asset, error) {
	return m.uploadFn(ctx, task)
}

type mockRemote struct {
	chat1.RemoteInterface
}

func (r mockRemote) GetS3Params(context.Context, chat1.ConversationID) (chat1.S3Params, error) {
	return chat1.S3Params{}, nil
}

func (r mockRemote) S3Sign(context.Context, chat1.S3SignArg) ([]byte, error) {
	return nil, nil
}

type mockActivityNotifier struct {
	types.ActivityNotifier
	startCh chan chat1.OutboxID
}

func newMockActivityNotifier() *mockActivityNotifier {
	return &mockActivityNotifier{
		startCh: make(chan chat1.OutboxID, 1000),
	}
}

func (a *mockActivityNotifier) AttachmentUploadStart(ctx context.Context, uid gregor1.UID,
	convID chat1.ConversationID, outboxID chat1.OutboxID) {
	a.startCh <- outboxID
}

func (a *mockActivityNotifier) AttachmentUploadProgress(ctx context.Context, uid gregor1.UID,
	convID chat1.ConversationID, outboxID chat1.OutboxID, bytesComplete, bytesTotal int64) {

}

type mockDeliverer struct {
	types.MessageDeliverer
	forceCh chan struct{}
}

func newMockDeliverer() *mockDeliverer {
	return &mockDeliverer{
		forceCh: make(chan struct{}, 1000),
	}
}

func (m *mockDeliverer) ForceDeliverLoop(context.Context) {
	m.forceCh <- struct{}{}
}

func (m *mockDeliverer) Stop(context.Context) chan struct{} {
	ch := make(chan struct{})
	close(ch)
	return ch
}

func TestAttachmentUploader(t *testing.T) {
	world := kbtest.NewChatMockWorld(t, "uploader", 1)
	defer world.Cleanup()

	u := world.GetUsers()[0]
	uid := gregor1.UID(u.User.GetUID().ToBytes())
	tc := world.Tcs[u.Username]
	g := globals.NewContext(tc.G, tc.ChatG)
	notifier := newMockActivityNotifier()
	store := &mockStore{}
	ri := mockRemote{}
	deliverer := newMockDeliverer()
	g.AttachmentURLSrv = types.DummyAttachmentHTTPSrv{}
	g.ActivityNotifier = notifier
	g.MessageDeliverer = deliverer
	getRi := func() chat1.RemoteInterface { return ri }
	cacheSize := 1
	uploader := NewUploader(g, store, NewS3Signer(getRi), getRi, cacheSize)
	convID := chat1.ConversationID([]byte{0, 1, 0})
	outboxID, err := storage.NewOutboxID()
	require.NoError(t, err)
	filename := "../testdata/ship.jpg"

	// Basic test to see if it works
	store.uploadFn = func(context.Context, *UploadTask) (chat1.Asset, error) {
		return chat1.Asset{}, nil
	}
	md, err := libkb.RandBytes(10)
	require.NoError(t, err)
	resChan, err := uploader.Register(context.TODO(), uid, convID, outboxID, "ship", filename, md, nil)
	require.NoError(t, err)
	uploadStartCheck := func(shouldHappen bool, outboxID chat1.OutboxID) {
		if shouldHappen {
			select {
			case obid := <-notifier.startCh:
				require.Equal(t, outboxID, obid)
			case <-time.After(20 * time.Second):
				require.Fail(t, "no start")
			}
		} else {
			select {
			case <-notifier.startCh:
				require.Fail(t, "start not supposed to happen")
			default:
			}
		}
	}
	deliverCheck := func(shouldHappen bool) {
		if shouldHappen {
			select {
			case <-deliverer.forceCh:
			case <-time.After(20 * time.Second):
				require.Fail(t, "no start")
			}
		} else {
			select {
			case <-deliverer.forceCh:
				require.Fail(t, "start not supposed to happen")
			default:
			}
		}
	}
	successCheck := func(cb types.AttachmentUploaderResultCb) {
		ch := cb.Wait()
		select {
		case res := <-ch:
			require.Nil(t, res.Error)
			require.Equal(t, md, res.Metadata)
			require.NotNil(t, res.Preview)
			require.Equal(t, "image/jpeg", res.Preview.MimeType)
			require.Equal(t, "image/jpeg", res.Object.MimeType)
		case <-time.After(20 * time.Second):
			require.Fail(t, "no upload")
		}
	}
	deliverCheck(true)
	uploadStartCheck(true, outboxID)
	successCheck(resChan)

	// Broken store
	outboxID, err = storage.NewOutboxID()
	require.NoError(t, err)
	store.uploadFn = func(context.Context, *UploadTask) (chat1.Asset, error) {
		return chat1.Asset{}, errors.New("i dont work")
	}
	resChan, err = uploader.Register(context.TODO(), uid, convID, outboxID, "ship", filename, md, nil)
	require.NoError(t, err)
	uploadStartCheck(true, outboxID)
	select {
	case res := <-resChan.Wait():
		require.NotNil(t, res.Error)
	case <-time.After(20 * time.Second):
		require.Fail(t, "no upload")
	}
	deliverCheck(true)

	// block until the upload is marked as done
	for count := 0; count <= 5; count++ {
		uploader.Lock()
		upload, ok := uploader.uploads[outboxID.String()]
		uploader.Unlock()
		if !ok && upload == nil {
			break
		}
		time.Sleep(time.Millisecond * 200)
		if count == 5 {
			require.Fail(t, "upload not marked as done")
		}
		t.Logf("upload not done, checking again")
	}
	t.Logf("upload done")

	// Retry after fixing store
	store.uploadFn = func(context.Context, *UploadTask) (chat1.Asset, error) {
		return chat1.Asset{}, nil
	}
	resChan, err = uploader.Retry(context.TODO(), outboxID)
	require.NoError(t, err)
	uploadStartCheck(true, outboxID)
	successCheck(resChan)
	deliverCheck(true)

	// Slow store to test concurrent retry
	outboxID, err = storage.NewOutboxID()
	require.NoError(t, err)
	slowCh := make(chan struct{})
	store.uploadFn = func(context.Context, *UploadTask) (chat1.Asset, error) {
		<-slowCh
		return chat1.Asset{}, nil
	}
	resChan, err = uploader.Register(context.TODO(), uid, convID, outboxID, "ship", filename, md, nil)
	require.NoError(t, err)
	uploadStartCheck(true, outboxID)
	deliverCheck(false)
	select {
	case <-resChan.Wait():
		require.Fail(t, "no res")
	default:
	}
	retryChan, err := uploader.Retry(context.TODO(), outboxID)
	require.NoError(t, err)
	uploadStartCheck(false, outboxID)
	close(slowCh)
	deliverCheck(true)
	// Should get results on both of these
	successCheck(retryChan)
	successCheck(resChan)

	uploader.Complete(context.TODO(), outboxID)
	_, _, err = uploader.Status(context.TODO(), outboxID)
	require.Error(t, err)

	// Test cancel
	outboxID, err = storage.NewOutboxID()
	require.NoError(t, err)
	slowCh = make(chan struct{})
	store.uploadFn = func(ctx context.Context, task *UploadTask) (chat1.Asset, error) {
		select {
		case <-slowCh:
		case <-ctx.Done():
			return chat1.Asset{}, ctx.Err()
		}
		return chat1.Asset{}, nil
	}
	resChan, err = uploader.Register(context.TODO(), uid, convID, outboxID, "ship", filename, md, nil)
	require.NoError(t, err)
	uploadStartCheck(true, outboxID)
	deliverCheck(false)
	select {
	case <-resChan.Wait():
		require.Fail(t, "no res")
	default:
	}
	require.NoError(t, uploader.Cancel(context.TODO(), outboxID))
	_, _, err = uploader.Status(context.TODO(), outboxID)
	require.Error(t, err)
	res := <-resChan.Wait()
	require.NotNil(t, res.Error)

	// verify uploadedPreviewsDir respects the cache size
	baseDir := uploader.getBaseDir()
	uploadedPreviews, err := filepath.Glob(filepath.Join(baseDir, uploadedPreviewsDir, "*"))
	require.NoError(t, err)
	require.Len(t, uploadedPreviews, 1)

	// verify uploadedFullsDir is respects the cache size
	uploadedFulls, err := filepath.Glob(filepath.Join(baseDir, uploadedFullsDir, "*"))
	require.NoError(t, err)
	require.Len(t, uploadedFulls, 1)
	mctx := kbtest.NewMetaContextForTest(*tc)

	// verify db nuke
	_, err = g.LocalDb.Nuke()
	require.NoError(t, err)
	err = uploader.OnDbNuke(mctx)
	require.NoError(t, err)

	uploadedPreviews, err = filepath.Glob(filepath.Join(baseDir, uploadedPreviewsDir, "*"))
	require.NoError(t, err)
	require.Zero(t, len(uploadedPreviews))

	uploadedFulls, err = filepath.Glob(filepath.Join(baseDir, uploadedFullsDir, "*"))
	require.NoError(t, err)
	require.Zero(t, len(uploadedFulls))
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// +build !production

package externals

import (
	"regexp"
	"strings"

	libkb "github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	jsonw "github.com/keybase/go-jsonw"
)

//=============================================================================
// Rooter
//

type RooterChecker struct {
	proof libkb.RemoteProofChainLink
}

var _ libkb.ProofChecker = (*RooterChecker)(nil)

func NewRooterChecker(p libkb.RemoteProofChainLink) (*RooterChecker, libkb.ProofError) {
	return &RooterChecker{p}, nil
}

func (rc *RooterChecker) GetTorError() libkb.ProofError { return nil }

func (rc *RooterChecker) CheckStatus(mctx libkb.MetaContext, h libkb.SigHint, _ libkb.ProofCheckerMode,
	pvlU keybase1.MerkleStoreEntry) (*libkb.SigHint, libkb.ProofError) {
	// TODO CORE-8951 see if we can populate verifiedHint with anything useful.
	return nil, CheckProofPvl(mctx, keybase1.ProofType_ROOTER, rc.proof, h, pvlU)
}

//
//=============================================================================

type RooterServiceType struct{ libkb.BaseServiceType }

func (t *RooterServiceType) Key() string { return t.GetTypeName() }

var rooterUsernameRegexp = regexp.MustCompile(`^(?i:[a-z0-9_]{1,20})$`)

func (t *RooterServiceType) NormalizeUsername(s string) (string, error) {
	if !rooterUsernameRegexp.MatchString(s) {
		return "", libkb.NewBadUsernameError(s)
	}
	return strings.ToLower(s), nil
}

func (t *RooterServiceType) NormalizeRemoteName(_ libkb.MetaContext, s string) (string, error) {
	// Allow a leading '@'.
	s = strings.TrimPrefix(s, "@")
	return t.NormalizeUsername(s)
}

func (t *RooterServiceType) GetPrompt() string {
	return "Your username on Rooter"
}

func (t *RooterServiceType) ToServiceJSON(un string) *jsonw.Wrapper {
	return t.BaseToServiceJSON(t, un)
}

func (t *RooterServiceType) PostInstructions(un string) *libkb.Markup {
	return libkb.FmtMarkup(`Please toot the following, and don't delete it:`)
}

func (t *RooterServiceType) DisplayName() string   { return "Rooter" }
func (t *RooterServiceType) GetTypeName() string   { return "rooter" }
func (t *RooterServiceType) PickerSubtext() string { return "" }
func (t *RooterServiceType) RecheckProofPosting(tryNumber int, status keybase1.ProofStatus, _ string) (warning *libkb.Markup, err error) {
	return t.BaseRecheckProofPosting(tryNumber, status)
}
func (t *RooterServiceType) GetProofType() string { return "test.web_service_binding.rooter" }

func (t *RooterServiceType) CheckProofText(text string, id keybase1.SigID, sig string) (err error) {
	return t.BaseCheckProofTextShort(text, id, true)
}

func (t *RooterServiceType) MakeProofChecker(l libkb.RemoteProofChainLink) libkb.ProofChecker {
	return &RooterChecker{l}
}

func (t *RooterServiceType) IsDevelOnly() bool { return true }
<commit_msg>hide rooter in prod mode (#16725)<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// +build !production

package externals

import (
	"regexp"
	"strings"

	libkb "github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	jsonw "github.com/keybase/go-jsonw"
)

//=============================================================================
// Rooter
//

type RooterChecker struct {
	proof libkb.RemoteProofChainLink
}

var _ libkb.ProofChecker = (*RooterChecker)(nil)

func NewRooterChecker(p libkb.RemoteProofChainLink) (*RooterChecker, libkb.ProofError) {
	return &RooterChecker{p}, nil
}

func (rc *RooterChecker) GetTorError() libkb.ProofError { return nil }

func (rc *RooterChecker) CheckStatus(mctx libkb.MetaContext, h libkb.SigHint, _ libkb.ProofCheckerMode,
	pvlU keybase1.MerkleStoreEntry) (*libkb.SigHint, libkb.ProofError) {
	// TODO CORE-8951 see if we can populate verifiedHint with anything useful.
	return nil, CheckProofPvl(mctx, keybase1.ProofType_ROOTER, rc.proof, h, pvlU)
}

//
//=============================================================================

type RooterServiceType struct{ libkb.BaseServiceType }

func (t *RooterServiceType) Key() string { return t.GetTypeName() }

var rooterUsernameRegexp = regexp.MustCompile(`^(?i:[a-z0-9_]{1,20})$`)

func (t *RooterServiceType) NormalizeUsername(s string) (string, error) {
	if !rooterUsernameRegexp.MatchString(s) {
		return "", libkb.NewBadUsernameError(s)
	}
	return strings.ToLower(s), nil
}

func (t *RooterServiceType) NormalizeRemoteName(_ libkb.MetaContext, s string) (string, error) {
	// Allow a leading '@'.
	s = strings.TrimPrefix(s, "@")
	return t.NormalizeUsername(s)
}

func (t *RooterServiceType) GetPrompt() string {
	return "Your username on Rooter"
}

func (t *RooterServiceType) CanMakeNewProofs(mctx libkb.MetaContext) bool {
	return mctx.G().GetRunMode() != libkb.ProductionRunMode
}

func (t *RooterServiceType) ToServiceJSON(un string) *jsonw.Wrapper {
	return t.BaseToServiceJSON(t, un)
}

func (t *RooterServiceType) PostInstructions(un string) *libkb.Markup {
	return libkb.FmtMarkup(`Please toot the following, and don't delete it:`)
}

func (t *RooterServiceType) DisplayName() string   { return "Rooter" }
func (t *RooterServiceType) GetTypeName() string   { return "rooter" }
func (t *RooterServiceType) PickerSubtext() string { return "" }
func (t *RooterServiceType) RecheckProofPosting(tryNumber int, status keybase1.ProofStatus, _ string) (warning *libkb.Markup, err error) {
	return t.BaseRecheckProofPosting(tryNumber, status)
}
func (t *RooterServiceType) GetProofType() string { return "test.web_service_binding.rooter" }

func (t *RooterServiceType) CheckProofText(text string, id keybase1.SigID, sig string) (err error) {
	return t.BaseCheckProofTextShort(text, id, true)
}

func (t *RooterServiceType) MakeProofChecker(l libkb.RemoteProofChainLink) libkb.ProofChecker {
	return &RooterChecker{l}
}

func (t *RooterServiceType) IsDevelOnly() bool { return true }
<|endoftext|>


<commit_before>package markdown

import (
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/mholt/caddy/middleware"
	"github.com/russross/blackfriday"
)

func TestMarkdown(t *testing.T) {
	templates := make(map[string]string)
	templates[DefaultTemplate] = "testdata/markdown_tpl.html"
	md := Markdown{
		Root:    "./testdata",
		FileSys: http.Dir("./testdata"),
		Configs: []Config{
			Config{
				Renderer:    blackfriday.HtmlRenderer(0, "", ""),
				PathScope:   "/blog",
				Extensions:  []string{".md"},
				Styles:      []string{},
				Scripts:     []string{},
				Templates:   templates,
				StaticDir:   DefaultStaticDir,
				StaticFiles: make(map[string]string),
			},
			Config{
				Renderer:    blackfriday.HtmlRenderer(0, "", ""),
				PathScope:   "/log",
				Extensions:  []string{".md"},
				Styles:      []string{"/resources/css/log.css", "/resources/css/default.css"},
				Scripts:     []string{"/resources/js/log.js", "/resources/js/default.js"},
				Templates:   make(map[string]string),
				StaticDir:   DefaultStaticDir,
				StaticFiles: make(map[string]string),
			},
			Config{
				Renderer:    blackfriday.HtmlRenderer(0, "", ""),
				PathScope:   "/og",
				Extensions:  []string{".md"},
				Styles:      []string{},
				Scripts:     []string{},
				Templates:   templates,
				StaticDir:   "testdata/og_static",
				StaticFiles: map[string]string{"/og/first.md": "testdata/og_static/og/first.md/index.html"},
				Links: []PageLink{
					PageLink{
						Title:   "first",
						Summary: "",
						Date:    time.Now(),
						Url:     "/og/first.md",
					},
				},
			},
		},
		IndexFiles: []string{"index.html"},
		Next: middleware.HandlerFunc(func(w http.ResponseWriter, r *http.Request) (int, error) {
			t.Fatalf("Next shouldn't be called")
			return 0, nil
		}),
	}

	req, err := http.NewRequest("GET", "/blog/test.md", nil)
	if err != nil {
		t.Fatalf("Could not create HTTP request: %v", err)
	}

	rec := httptest.NewRecorder()

	md.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("Wrong status, expected: %d and got %d", http.StatusOK, rec.Code)
	}

	respBody := rec.Body.String()
	expectedBody := `<!DOCTYPE html>

<title>Markdown test</title>

Header

Welcome to A Caddy website!

Welcome on the blog

Body

go
func getTrue() bool {
    return true
}

`
	if respBody != expectedBody {
		t.Fatalf("Expected body: %v got: %v", expectedBody, respBody)
	}

	req, err = http.NewRequest("GET", "/log/test.md", nil)
	if err != nil {
		t.Fatalf("Could not create HTTP request: %v", err)
	}
	rec = httptest.NewRecorder()

	md.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("Wrong status, expected: %d and got %d", http.StatusOK, rec.Code)
	}
	respBody = rec.Body.String()
	expectedBody = `<!DOCTYPE html>

		<title>Markdown test</title>
		<meta charset="utf-8">
		<link rel="stylesheet" href="/resources/css/log.css">
<link rel="stylesheet" href="/resources/css/default.css">

Welcome on the blog

Body

go
func getTrue() bool {
    return true
}

`

	replacer := strings.NewReplacer("\r", "", "\n", "")
	respBody = replacer.Replace(respBody)
	expectedBody = replacer.Replace(expectedBody)
	if respBody != expectedBody {
		t.Fatalf("Expected body: %v got: %v", expectedBody, respBody)
	}

	req, err = http.NewRequest("GET", "/og/first.md", nil)
	if err != nil {
		t.Fatalf("Could not create HTTP request: %v", err)
	}
	rec = httptest.NewRecorder()
	currenttime := time.Now().Local().Add(-time.Second)
	err = os.Chtimes("testdata/og/first.md", currenttime, currenttime)
	currenttime = time.Now().Local()
	err = os.Chtimes("testdata/og_static/og/first.md/index.html", currenttime, currenttime)

	md.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("Wrong status, expected: %d and got %d", http.StatusOK, rec.Code)
	}
	respBody = rec.Body.String()
	expectedBody = `<!DOCTYPE html>

<title>first_post</title>

Header title

Test h1

`
	respBody = replacer.Replace(respBody)
	expectedBody = replacer.Replace(expectedBody)
	if respBody != expectedBody {
		t.Fatalf("Expected body: %v got: %v", expectedBody, respBody)
	}

	expectedLinks := []string{
		"/blog/test.md",
		"/log/test.md",
		"/og/first.md",
	}

	for i, c := range md.Configs {
		log.Printf("Test number: %d, configuration links: %v, config: %v", i, c.Links, c)
		if c.Links[0].Url != expectedLinks[i] {
			t.Fatalf("Expected %v got %v", expectedLinks[i], c.Links[0].Url)
		}
	}

	// attempt to trigger race condition
	var w sync.WaitGroup
	f := func() {
		req, err := http.NewRequest("GET", "/log/test.md", nil)
		if err != nil {
			t.Fatalf("Could not create HTTP request: %v", err)
		}
		rec := httptest.NewRecorder()

		md.ServeHTTP(rec, req)
		w.Done()
	}
	for i := 0; i < 5; i++ {
		w.Add(1)
		go f()
	}
	w.Wait()

	if err = os.RemoveAll(DefaultStaticDir); err != nil {
		t.Errorf("Error while removing the generated static files: %v", err)
	}

}
<commit_msg>Fix markdown tests that I broke<commit_after>package markdown

import (
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/mholt/caddy/middleware"
	"github.com/russross/blackfriday"
)

func TestMarkdown(t *testing.T) {
	templates := make(map[string]string)
	templates[DefaultTemplate] = "testdata/markdown_tpl.html"
	md := Markdown{
		Root:    "./testdata",
		FileSys: http.Dir("./testdata"),
		Configs: []Config{
			Config{
				Renderer:    blackfriday.HtmlRenderer(0, "", ""),
				PathScope:   "/blog",
				Extensions:  []string{".md"},
				Styles:      []string{},
				Scripts:     []string{},
				Templates:   templates,
				StaticDir:   DefaultStaticDir,
				StaticFiles: make(map[string]string),
			},
			Config{
				Renderer:    blackfriday.HtmlRenderer(0, "", ""),
				PathScope:   "/log",
				Extensions:  []string{".md"},
				Styles:      []string{"/resources/css/log.css", "/resources/css/default.css"},
				Scripts:     []string{"/resources/js/log.js", "/resources/js/default.js"},
				Templates:   make(map[string]string),
				StaticDir:   DefaultStaticDir,
				StaticFiles: make(map[string]string),
			},
			Config{
				Renderer:    blackfriday.HtmlRenderer(0, "", ""),
				PathScope:   "/og",
				Extensions:  []string{".md"},
				Styles:      []string{},
				Scripts:     []string{},
				Templates:   templates,
				StaticDir:   "testdata/og_static",
				StaticFiles: map[string]string{"/og/first.md": "testdata/og_static/og/first.md/index.html"},
				Links: []PageLink{
					PageLink{
						Title:   "first",
						Summary: "",
						Date:    time.Now(),
						URL:     "/og/first.md",
					},
				},
			},
		},
		IndexFiles: []string{"index.html"},
		Next: middleware.HandlerFunc(func(w http.ResponseWriter, r *http.Request) (int, error) {
			t.Fatalf("Next shouldn't be called")
			return 0, nil
		}),
	}

	req, err := http.NewRequest("GET", "/blog/test.md", nil)
	if err != nil {
		t.Fatalf("Could not create HTTP request: %v", err)
	}

	rec := httptest.NewRecorder()

	md.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("Wrong status, expected: %d and got %d", http.StatusOK, rec.Code)
	}

	respBody := rec.Body.String()
	expectedBody := `<!DOCTYPE html>

<title>Markdown test</title>

Header

Welcome to A Caddy website!

Welcome on the blog

Body

go
func getTrue() bool {
    return true
}

`
	if respBody != expectedBody {
		t.Fatalf("Expected body: %v got: %v", expectedBody, respBody)
	}

	req, err = http.NewRequest("GET", "/log/test.md", nil)
	if err != nil {
		t.Fatalf("Could not create HTTP request: %v", err)
	}
	rec = httptest.NewRecorder()

	md.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("Wrong status, expected: %d and got %d", http.StatusOK, rec.Code)
	}
	respBody = rec.Body.String()
	expectedBody = `<!DOCTYPE html>

		<title>Markdown test</title>
		<meta charset="utf-8">
		<link rel="stylesheet" href="/resources/css/log.css">
<link rel="stylesheet" href="/resources/css/default.css">

Welcome on the blog

Body

go
func getTrue() bool {
    return true
}

`

	replacer := strings.NewReplacer("\r", "", "\n", "")
	respBody = replacer.Replace(respBody)
	expectedBody = replacer.Replace(expectedBody)
	if respBody != expectedBody {
		t.Fatalf("Expected body: %v got: %v", expectedBody, respBody)
	}

	req, err = http.NewRequest("GET", "/og/first.md", nil)
	if err != nil {
		t.Fatalf("Could not create HTTP request: %v", err)
	}
	rec = httptest.NewRecorder()
	currenttime := time.Now().Local().Add(-time.Second)
	err = os.Chtimes("testdata/og/first.md", currenttime, currenttime)
	currenttime = time.Now().Local()
	err = os.Chtimes("testdata/og_static/og/first.md/index.html", currenttime, currenttime)

	md.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("Wrong status, expected: %d and got %d", http.StatusOK, rec.Code)
	}
	respBody = rec.Body.String()
	expectedBody = `<!DOCTYPE html>

<title>first_post</title>

Header title

Test h1

`
	respBody = replacer.Replace(respBody)
	expectedBody = replacer.Replace(expectedBody)
	if respBody != expectedBody {
		t.Fatalf("Expected body: %v got: %v", expectedBody, respBody)
	}

	expectedLinks := []string{
		"/blog/test.md",
		"/log/test.md",
		"/og/first.md",
	}

	for i, c := range md.Configs {
		log.Printf("Test number: %d, configuration links: %v, config: %v", i, c.Links, c)
		if c.Links[0].URL != expectedLinks[i] {
			t.Fatalf("Expected %v got %v", expectedLinks[i], c.Links[0].URL)
		}
	}

	// attempt to trigger race condition
	var w sync.WaitGroup
	f := func() {
		req, err := http.NewRequest("GET", "/log/test.md", nil)
		if err != nil {
			t.Fatalf("Could not create HTTP request: %v", err)
		}
		rec := httptest.NewRecorder()

		md.ServeHTTP(rec, req)
		w.Done()
	}
	for i := 0; i < 5; i++ {
		w.Add(1)
		go f()
	}
	w.Wait()

	if err = os.RemoveAll(DefaultStaticDir); err != nil {
		t.Errorf("Error while removing the generated static files: %v", err)
	}

}
<|endoftext|>


<commit_before>package gitlabnet

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"path"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet/testserver"
	"gitlab.com/gitlab-org/gitlab-shell/go/internal/testhelper"
)

func TestClients(t *testing.T) {
	testDirCleanup, err := testhelper.PrepareTestRootDir()
	require.NoError(t, err)
	defer testDirCleanup()

	requests := []testserver.TestRequestHandler{
		{
			Path: "/api/v4/internal/hello",
			Handler: func(w http.ResponseWriter, r *http.Request) {
				require.Equal(t, http.MethodGet, r.Method)

				fmt.Fprint(w, "Hello")
			},
		},
		{
			Path: "/api/v4/internal/post_endpoint",
			Handler: func(w http.ResponseWriter, r *http.Request) {
				require.Equal(t, http.MethodPost, r.Method)

				b, err := ioutil.ReadAll(r.Body)
				defer r.Body.Close()

				require.NoError(t, err)

				fmt.Fprint(w, "Echo: "+string(b))
			},
		},
		{
			Path: "/api/v4/internal/auth",
			Handler: func(w http.ResponseWriter, r *http.Request) {
				fmt.Fprint(w, r.Header.Get(secretHeaderName))
			},
		},
		{
			Path: "/api/v4/internal/error",
			Handler: func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "application/json")
				w.WriteHeader(http.StatusBadRequest)
				body := map[string]string{
					"message": "Don't do that",
				}
				json.NewEncoder(w).Encode(body)
			},
		},
		{
			Path: "/api/v4/internal/broken",
			Handler: func(w http.ResponseWriter, r *http.Request) {
				panic("Broken")
			},
		},
	}

	testCases := []struct {
		desc   string
		config *config.Config
		server func(*testing.T, []testserver.TestRequestHandler) (string, func())
	}{
		{
			desc:   "Socket client",
			config: &config.Config{},
			server: testserver.StartSocketHttpServer,
		},
		{
			desc:   "Http client",
			config: &config.Config{},
			server: testserver.StartHttpServer,
		},
		{
			desc: "Https client",
			config: &config.Config{
				HttpSettings: config.HttpSettingsConfig{CaFile: path.Join(testhelper.TestRoot, "certs/valid/server.crt")},
			},
			server: testserver.StartHttpsServer,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.desc, func(t *testing.T) {
			url, cleanup := tc.server(t, requests)
			defer cleanup()

			tc.config.GitlabUrl = url
			tc.config.Secret = "sssh, it's a secret"

			client, err := GetClient(tc.config)
			require.NoError(t, err)

			testBrokenRequest(t, client)
			testSuccessfulGet(t, client)
			testSuccessfulPost(t, client)
			testMissing(t, client)
			testErrorMessage(t, client)
			testAuthenticationHeader(t, client)
		})
	}
}

func testSuccessfulGet(t *testing.T, client *GitlabClient) {
	t.Run("Successful get", func(t *testing.T) {
		response, err := client.Get("/hello")
		defer response.Body.Close()

		require.NoError(t, err)
		require.NotNil(t, response)

		responseBody, err := ioutil.ReadAll(response.Body)
		assert.NoError(t, err)
		assert.Equal(t, string(responseBody), "Hello")
	})
}

func testSuccessfulPost(t *testing.T, client *GitlabClient) {
	t.Run("Successful Post", func(t *testing.T) {
		data := map[string]string{"key": "value"}

		response, err := client.Post("/post_endpoint", data)
		defer response.Body.Close()

		require.NoError(t, err)
		require.NotNil(t, response)

		responseBody, err := ioutil.ReadAll(response.Body)
		assert.NoError(t, err)
		assert.Equal(t, "Echo: {\"key\":\"value\"}", string(responseBody))
	})
}

func testMissing(t *testing.T, client *GitlabClient) {
	t.Run("Missing error for GET", func(t *testing.T) {
		response, err := client.Get("/missing")
		assert.EqualError(t, err, "Internal API error (404)")
		assert.Nil(t, response)
	})

	t.Run("Missing error for POST", func(t *testing.T) {
		response, err := client.Post("/missing", map[string]string{})
		assert.EqualError(t, err, "Internal API error (404)")
		assert.Nil(t, response)
	})
}

func testErrorMessage(t *testing.T, client *GitlabClient) {
	t.Run("Error with message for GET", func(t *testing.T) {
		response, err := client.Get("/error")
		assert.EqualError(t, err, "Don't do that")
		assert.Nil(t, response)
	})

	t.Run("Error with message for POST", func(t *testing.T) {
		response, err := client.Post("/error", map[string]string{})
		assert.EqualError(t, err, "Don't do that")
		assert.Nil(t, response)
	})
}

func testBrokenRequest(t *testing.T, client *GitlabClient) {
	t.Run("Broken request for GET", func(t *testing.T) {
		response, err := client.Get("/broken")
		assert.EqualError(t, err, "Internal API unreachable")
		assert.Nil(t, response)
	})

	t.Run("Broken request for POST", func(t *testing.T) {
		response, err := client.Post("/broken", map[string]string{})
		assert.EqualError(t, err, "Internal API unreachable")
		assert.Nil(t, response)
	})
}

func testAuthenticationHeader(t *testing.T, client *GitlabClient) {
	t.Run("Authentication headers for GET", func(t *testing.T) {
		response, err := client.Get("/auth")
		defer response.Body.Close()

		require.NoError(t, err)
		require.NotNil(t, response)

		responseBody, err := ioutil.ReadAll(response.Body)
		require.NoError(t, err)

		header, err := base64.StdEncoding.DecodeString(string(responseBody))
		require.NoError(t, err)
		assert.Equal(t, "sssh, it's a secret", string(header))
	})

	t.Run("Authentication headers for POST", func(t *testing.T) {
		response, err := client.Post("/auth", map[string]string{})
		defer response.Body.Close()

		require.NoError(t, err)
		require.NotNil(t, response)

		responseBody, err := ioutil.ReadAll(response.Body)
		require.NoError(t, err)

		header, err := base64.StdEncoding.DecodeString(string(responseBody))
		require.NoError(t, err)
		assert.Equal(t, "sssh, it's a secret", string(header))
	})
}
<commit_msg>Fix logic errors in gitlabnet client tests<commit_after>package gitlabnet

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"path"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet/testserver"
	"gitlab.com/gitlab-org/gitlab-shell/go/internal/testhelper"
)

func TestClients(t *testing.T) {
	testDirCleanup, err := testhelper.PrepareTestRootDir()
	require.NoError(t, err)
	defer testDirCleanup()

	requests := []testserver.TestRequestHandler{
		{
			Path: "/api/v4/internal/hello",
			Handler: func(w http.ResponseWriter, r *http.Request) {
				require.Equal(t, http.MethodGet, r.Method)

				fmt.Fprint(w, "Hello")
			},
		},
		{
			Path: "/api/v4/internal/post_endpoint",
			Handler: func(w http.ResponseWriter, r *http.Request) {
				require.Equal(t, http.MethodPost, r.Method)

				b, err := ioutil.ReadAll(r.Body)
				defer r.Body.Close()

				require.NoError(t, err)

				fmt.Fprint(w, "Echo: "+string(b))
			},
		},
		{
			Path: "/api/v4/internal/auth",
			Handler: func(w http.ResponseWriter, r *http.Request) {
				fmt.Fprint(w, r.Header.Get(secretHeaderName))
			},
		},
		{
			Path: "/api/v4/internal/error",
			Handler: func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "application/json")
				w.WriteHeader(http.StatusBadRequest)
				body := map[string]string{
					"message": "Don't do that",
				}
				json.NewEncoder(w).Encode(body)
			},
		},
		{
			Path: "/api/v4/internal/broken",
			Handler: func(w http.ResponseWriter, r *http.Request) {
				panic("Broken")
			},
		},
	}

	testCases := []struct {
		desc   string
		config *config.Config
		server func(*testing.T, []testserver.TestRequestHandler) (string, func())
	}{
		{
			desc:   "Socket client",
			config: &config.Config{},
			server: testserver.StartSocketHttpServer,
		},
		{
			desc:   "Http client",
			config: &config.Config{},
			server: testserver.StartHttpServer,
		},
		{
			desc: "Https client",
			config: &config.Config{
				HttpSettings: config.HttpSettingsConfig{CaFile: path.Join(testhelper.TestRoot, "certs/valid/server.crt")},
			},
			server: testserver.StartHttpsServer,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.desc, func(t *testing.T) {
			url, cleanup := tc.server(t, requests)
			defer cleanup()

			tc.config.GitlabUrl = url
			tc.config.Secret = "sssh, it's a secret"

			client, err := GetClient(tc.config)
			require.NoError(t, err)

			testBrokenRequest(t, client)
			testSuccessfulGet(t, client)
			testSuccessfulPost(t, client)
			testMissing(t, client)
			testErrorMessage(t, client)
			testAuthenticationHeader(t, client)
		})
	}
}

func testSuccessfulGet(t *testing.T, client *GitlabClient) {
	t.Run("Successful get", func(t *testing.T) {
		response, err := client.Get("/hello")
		require.NoError(t, err)
		require.NotNil(t, response)

		defer response.Body.Close()

		responseBody, err := ioutil.ReadAll(response.Body)
		assert.NoError(t, err)
		assert.Equal(t, string(responseBody), "Hello")
	})
}

func testSuccessfulPost(t *testing.T, client *GitlabClient) {
	t.Run("Successful Post", func(t *testing.T) {
		data := map[string]string{"key": "value"}

		response, err := client.Post("/post_endpoint", data)
		require.NoError(t, err)
		require.NotNil(t, response)

		defer response.Body.Close()

		responseBody, err := ioutil.ReadAll(response.Body)
		assert.NoError(t, err)
		assert.Equal(t, "Echo: {\"key\":\"value\"}", string(responseBody))
	})
}

func testMissing(t *testing.T, client *GitlabClient) {
	t.Run("Missing error for GET", func(t *testing.T) {
		response, err := client.Get("/missing")
		assert.EqualError(t, err, "Internal API error (404)")
		assert.Nil(t, response)
	})

	t.Run("Missing error for POST", func(t *testing.T) {
		response, err := client.Post("/missing", map[string]string{})
		assert.EqualError(t, err, "Internal API error (404)")
		assert.Nil(t, response)
	})
}

func testErrorMessage(t *testing.T, client *GitlabClient) {
	t.Run("Error with message for GET", func(t *testing.T) {
		response, err := client.Get("/error")
		assert.EqualError(t, err, "Don't do that")
		assert.Nil(t, response)
	})

	t.Run("Error with message for POST", func(t *testing.T) {
		response, err := client.Post("/error", map[string]string{})
		assert.EqualError(t, err, "Don't do that")
		assert.Nil(t, response)
	})
}

func testBrokenRequest(t *testing.T, client *GitlabClient) {
	t.Run("Broken request for GET", func(t *testing.T) {
		response, err := client.Get("/broken")
		assert.EqualError(t, err, "Internal API unreachable")
		assert.Nil(t, response)
	})

	t.Run("Broken request for POST", func(t *testing.T) {
		response, err := client.Post("/broken", map[string]string{})
		assert.EqualError(t, err, "Internal API unreachable")
		assert.Nil(t, response)
	})
}

func testAuthenticationHeader(t *testing.T, client *GitlabClient) {
	t.Run("Authentication headers for GET", func(t *testing.T) {
		response, err := client.Get("/auth")
		require.NoError(t, err)
		require.NotNil(t, response)

		defer response.Body.Close()

		responseBody, err := ioutil.ReadAll(response.Body)
		require.NoError(t, err)

		header, err := base64.StdEncoding.DecodeString(string(responseBody))
		require.NoError(t, err)
		assert.Equal(t, "sssh, it's a secret", string(header))
	})

	t.Run("Authentication headers for POST", func(t *testing.T) {
		response, err := client.Post("/auth", map[string]string{})
		require.NoError(t, err)
		require.NotNil(t, response)

		defer response.Body.Close()

		responseBody, err := ioutil.ReadAll(response.Body)
		require.NoError(t, err)

		header, err := base64.StdEncoding.DecodeString(string(responseBody))
		require.NoError(t, err)
		assert.Equal(t, "sssh, it's a secret", string(header))
	})
}
<|endoftext|>


<commit_before>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package sqlparser

import (
	"testing"

	"github.com/youtube/vitess/go/sqltypes"
)

func TestParsedQuery(t *testing.T) {
	tcases := []struct {
		desc     string
		query    string
		bindVars map[string]interface{}
		output   string
	}{
		{
			"no subs",
			"select * from a where id = 2",
			map[string]interface{}{
				"id": 1,
			},
			"select * from a where id = 2",
		}, {
			"simple bindvar sub",
			"select * from a where id1 = :id1 and id2 = :id2",
			map[string]interface{}{
				"id1": 1,
				"id2": nil,
			},
			"select * from a where id1 = 1 and id2 = null",
		}, {
			"missing bind var",
			"select * from a where id1 = :id1 and id2 = :id2",
			map[string]interface{}{
				"id1": 1,
			},
			"missing bind var id2",
		}, {
			"unencodable bind var",
			"select * from a where id1 = :id",
			map[string]interface{}{
				"id": make([]int, 1),
			},
			"unsupported bind variable type []int: [0]",
		}, {
			"list inside bind vars",
			"select * from a where id in (:vals)",
			map[string]interface{}{
				"vals": []sqltypes.Value{
					sqltypes.MakeNumeric([]byte("1")),
					sqltypes.MakeString([]byte("aa")),
				},
			},
			"select * from a where id in (1, 'aa')",
		}, {
			"two lists inside bind vars",
			"select * from a where id in (:vals)",
			map[string]interface{}{
				"vals": [][]sqltypes.Value{
					[]sqltypes.Value{
						sqltypes.MakeNumeric([]byte("1")),
						sqltypes.MakeString([]byte("aa")),
					},
					[]sqltypes.Value{
						sqltypes.Value{},
						sqltypes.MakeString([]byte("bb")),
					},
				},
			},
			"select * from a where id in ((1, 'aa'), (null, 'bb'))",
		}, {
			"list bind vars",
			"select * from a where id in ::vals",
			map[string]interface{}{
				"vals": []interface{}{
					1,
					"aa",
				},
			},
			"select * from a where id in (1, 'aa')",
		}, {
			"list bind vars single argument",
			"select * from a where id in ::vals",
			map[string]interface{}{
				"vals": []interface{}{
					1,
				},
			},
			"select * from a where id in (1)",
		}, {
			"list bind vars 0 arguments",
			"select * from a where id in ::vals",
			map[string]interface{}{
				"vals": []interface{}{},
			},
			"empty list supplied for vals",
		}, {
			"non-list bind var supplied",
			"select * from a where id in ::vals",
			map[string]interface{}{
				"vals": 1,
			},
			"unexpected list arg type int for key vals",
		}, {
			"list bind var for non-list",
			"select * from a where id = :vals",
			map[string]interface{}{
				"vals": []interface{}{1},
			},
			"unexpected arg type []interface {} for key vals",
		}, {
			"single column tuple equality",
			// We have to use an incorrect construct to get around the parser.
			"select * from a where b = :equality",
			map[string]interface{}{
				"equality": TupleEqualityList{
					Columns: []string{"pk"},
					Rows: [][]sqltypes.Value{
						[]sqltypes.Value{sqltypes.MakeNumeric([]byte("1"))},
						[]sqltypes.Value{sqltypes.MakeString([]byte("aa"))},
					},
				},
			},
			"select * from a where b = pk in (1, 'aa')",
		}, {
			"multi column tuple equality",
			"select * from a where b = :equality",
			map[string]interface{}{
				"equality": TupleEqualityList{
					Columns: []string{"pk1", "pk2"},
					Rows: [][]sqltypes.Value{
						[]sqltypes.Value{
							sqltypes.MakeNumeric([]byte("1")),
							sqltypes.MakeString([]byte("aa")),
						},
						[]sqltypes.Value{
							sqltypes.MakeNumeric([]byte("2")),
							sqltypes.MakeString([]byte("bb")),
						},
					},
				},
			},
			"select * from a where b = (pk1 = 1 and pk2 = 'aa') or (pk1 = 2 and pk2 = 'bb')",
		}, {
			"0 rows",
			"select * from a where b = :equality",
			map[string]interface{}{
				"equality": TupleEqualityList{
					Columns: []string{"pk"},
					Rows:    [][]sqltypes.Value{},
				},
			},
			"cannot encode with 0 rows",
		}, {
			"values don't match column count",
			"select * from a where b = :equality",
			map[string]interface{}{
				"equality": TupleEqualityList{
					Columns: []string{"pk"},
					Rows: [][]sqltypes.Value{
						[]sqltypes.Value{
							sqltypes.MakeNumeric([]byte("1")),
							sqltypes.MakeString([]byte("aa")),
						},
					},
				},
			},
			"values don't match column count",
		},
	}

	for _, tcase := range tcases {
		tree, err := Parse(tcase.query)
		if err != nil {
			t.Errorf("parse failed for %s: %v", tcase.desc, err)
			continue
		}
		buf := NewTrackedBuffer(nil)
		buf.Myprintf("%v", tree)
		pq := buf.ParsedQuery()
		bytes, err := pq.GenerateQuery(tcase.bindVars)
		var got string
		if err != nil {
			got = err.Error()
		} else {
			got = string(bytes)
		}
		if got != tcase.output {
			t.Errorf("for test case: %s, got: '%s', want '%s'", tcase.desc, got, tcase.output)
		}
	}
}
<commit_msg>sqlparser: add test for coverage<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package sqlparser

import (
	"reflect"
	"testing"

	"github.com/youtube/vitess/go/sqltypes"
)

func TestParsedQuery(t *testing.T) {
	tcases := []struct {
		desc     string
		query    string
		bindVars map[string]interface{}
		output   string
	}{
		{
			"no subs",
			"select * from a where id = 2",
			map[string]interface{}{
				"id": 1,
			},
			"select * from a where id = 2",
		}, {
			"simple bindvar sub",
			"select * from a where id1 = :id1 and id2 = :id2",
			map[string]interface{}{
				"id1": 1,
				"id2": nil,
			},
			"select * from a where id1 = 1 and id2 = null",
		}, {
			"missing bind var",
			"select * from a where id1 = :id1 and id2 = :id2",
			map[string]interface{}{
				"id1": 1,
			},
			"missing bind var id2",
		}, {
			"unencodable bind var",
			"select * from a where id1 = :id",
			map[string]interface{}{
				"id": make([]int, 1),
			},
			"unsupported bind variable type []int: [0]",
		}, {
			"list inside bind vars",
			"select * from a where id in (:vals)",
			map[string]interface{}{
				"vals": []sqltypes.Value{
					sqltypes.MakeNumeric([]byte("1")),
					sqltypes.MakeString([]byte("aa")),
				},
			},
			"select * from a where id in (1, 'aa')",
		}, {
			"two lists inside bind vars",
			"select * from a where id in (:vals)",
			map[string]interface{}{
				"vals": [][]sqltypes.Value{
					[]sqltypes.Value{
						sqltypes.MakeNumeric([]byte("1")),
						sqltypes.MakeString([]byte("aa")),
					},
					[]sqltypes.Value{
						sqltypes.Value{},
						sqltypes.MakeString([]byte("bb")),
					},
				},
			},
			"select * from a where id in ((1, 'aa'), (null, 'bb'))",
		}, {
			"list bind vars",
			"select * from a where id in ::vals",
			map[string]interface{}{
				"vals": []interface{}{
					1,
					"aa",
				},
			},
			"select * from a where id in (1, 'aa')",
		}, {
			"list bind vars single argument",
			"select * from a where id in ::vals",
			map[string]interface{}{
				"vals": []interface{}{
					1,
				},
			},
			"select * from a where id in (1)",
		}, {
			"list bind vars 0 arguments",
			"select * from a where id in ::vals",
			map[string]interface{}{
				"vals": []interface{}{},
			},
			"empty list supplied for vals",
		}, {
			"non-list bind var supplied",
			"select * from a where id in ::vals",
			map[string]interface{}{
				"vals": 1,
			},
			"unexpected list arg type int for key vals",
		}, {
			"list bind var for non-list",
			"select * from a where id = :vals",
			map[string]interface{}{
				"vals": []interface{}{1},
			},
			"unexpected arg type []interface {} for key vals",
		}, {
			"single column tuple equality",
			// We have to use an incorrect construct to get around the parser.
			"select * from a where b = :equality",
			map[string]interface{}{
				"equality": TupleEqualityList{
					Columns: []string{"pk"},
					Rows: [][]sqltypes.Value{
						[]sqltypes.Value{sqltypes.MakeNumeric([]byte("1"))},
						[]sqltypes.Value{sqltypes.MakeString([]byte("aa"))},
					},
				},
			},
			"select * from a where b = pk in (1, 'aa')",
		}, {
			"multi column tuple equality",
			"select * from a where b = :equality",
			map[string]interface{}{
				"equality": TupleEqualityList{
					Columns: []string{"pk1", "pk2"},
					Rows: [][]sqltypes.Value{
						[]sqltypes.Value{
							sqltypes.MakeNumeric([]byte("1")),
							sqltypes.MakeString([]byte("aa")),
						},
						[]sqltypes.Value{
							sqltypes.MakeNumeric([]byte("2")),
							sqltypes.MakeString([]byte("bb")),
						},
					},
				},
			},
			"select * from a where b = (pk1 = 1 and pk2 = 'aa') or (pk1 = 2 and pk2 = 'bb')",
		}, {
			"0 rows",
			"select * from a where b = :equality",
			map[string]interface{}{
				"equality": TupleEqualityList{
					Columns: []string{"pk"},
					Rows:    [][]sqltypes.Value{},
				},
			},
			"cannot encode with 0 rows",
		}, {
			"values don't match column count",
			"select * from a where b = :equality",
			map[string]interface{}{
				"equality": TupleEqualityList{
					Columns: []string{"pk"},
					Rows: [][]sqltypes.Value{
						[]sqltypes.Value{
							sqltypes.MakeNumeric([]byte("1")),
							sqltypes.MakeString([]byte("aa")),
						},
					},
				},
			},
			"values don't match column count",
		},
	}

	for _, tcase := range tcases {
		tree, err := Parse(tcase.query)
		if err != nil {
			t.Errorf("parse failed for %s: %v", tcase.desc, err)
			continue
		}
		buf := NewTrackedBuffer(nil)
		buf.Myprintf("%v", tree)
		pq := buf.ParsedQuery()
		bytes, err := pq.GenerateQuery(tcase.bindVars)
		var got string
		if err != nil {
			got = err.Error()
		} else {
			got = string(bytes)
		}
		if got != tcase.output {
			t.Errorf("for test case: %s, got: '%s', want '%s'", tcase.desc, got, tcase.output)
		}
	}
}

func TestGenerateParsedQuery(t *testing.T) {
	stmt, err := Parse("select * from a where id =:id")
	if err != nil {
		t.Error(err)
		return
	}
	pq := GenerateParsedQuery(stmt)
	want := &ParsedQuery{
		Query:         "select * from a where id = :id",
		bindLocations: []bindLocation{{offset: 27, length: 3}},
	}
	if !reflect.DeepEqual(pq, want) {
		t.Errorf("GenerateParsedQuery: %+v, want %+v", pq, want)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Vitess Authors.

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

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

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

package vindexes

import (
	"encoding/binary"
	"encoding/json"
	"io/ioutil"

	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/key"
	"vitess.io/vitess/go/vt/log"
)

var (
	_ MultiColumn = (*RegionJson)(nil)
)

func init() {
	Register("region_json", NewRegionJson)
}

// RegionMap is used to store mapping of country to region
type RegionMap map[string]uint64

// RegionJson defines a vindex that uses a lookup table.
// The table is expected to define the id column as unique. It's
// Unique and a Lookup.
type RegionJson struct {
	name        string
	regionMap   RegionMap
	regionBytes int
}

// NewRegionJson creates a RegionJson vindex.
// The supplied map requires all the fields of "RegionExperimental".
// Additionally, it requires a region_map argument representing the path to a json file
// containing a map of country to region.
func NewRegionJson(name string, m map[string]string) (Vindex, error) {
	rmPath := m["region_map"]
	rmap := make(map[string]uint64)
	data, err := ioutil.ReadFile(rmPath)
	if err != nil {
		return nil, err
	}
	log.Infof("Loaded Region map from: %s", rmPath)
	err = json.Unmarshal(data, &rmap)
	if err != nil {
		return nil, err
	}

	return &RegionJson{
		name:      name,
		regionMap: rmap,
	}, nil
}

// String returns the name of the vindex.
func (rv *RegionJson) String() string {
	return rv.name
}

// Cost returns the cost of this index as 1.
func (rv *RegionJson) Cost() int {
	return 1
}

// IsUnique returns true since the Vindex is unique.
func (rv *RegionJson) IsUnique() bool {
	return true
}

// Map satisfies MultiColumn.
func (rv *RegionJson) Map(vcursor VCursor, rowsColValues [][]sqltypes.Value) ([]key.Destination, error) {
	destinations := make([]key.Destination, 0, len(rowsColValues))
	for _, row := range rowsColValues {
		if len(row) != 2 {
			destinations = append(destinations, key.DestinationNone{})
			continue
		}
		// Compute hash.
		hn, err := sqltypes.ToUint64(row[0])
		if err != nil {
			destinations = append(destinations, key.DestinationNone{})
			continue
		}
		h := vhash(hn)

		rn, ok := rv.regionMap[row[1].ToString()]
		if !ok {
			destinations = append(destinations, key.DestinationNone{})
			continue
		}
		r := make([]byte, 2)
		binary.BigEndian.PutUint16(r, uint16(rn))

		// Concatenate and add to destinations.
		if rv.regionBytes == 1 {
			r = r[1:]
		}
		dest := append(r, h...)
		destinations = append(destinations, key.DestinationKeyspaceID(dest))
	}
	return destinations, nil
}

// Verify satisfies MultiColumn
func (rv *RegionJson) Verify(vcursor VCursor, rowsColValues [][]sqltypes.Value, ksids [][]byte) ([]bool, error) {
	result := make([]bool, len(rowsColValues))
	destinations, _ := rv.Map(vcursor, rowsColValues)
	for i, dest := range destinations {
		destksid, ok := dest.(key.DestinationKeyspaceID)
		if !ok {
			continue
		}
		result[i] = bytes.Equal([]byte(destksid), ksids[i])
	}
	return result, nil
}

// NeedVCursor satisfies the Vindex interface.
func (rv *RegionJson) NeedsVCursor() bool {
	return false
}
<commit_msg>Added bytes back in.<commit_after>/*
Copyright 2020 The Vitess Authors.

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

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

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

package vindexes

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"io/ioutil"

	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/key"
	"vitess.io/vitess/go/vt/log"
)

var (
	_ MultiColumn = (*RegionJson)(nil)
)

func init() {
	Register("region_json", NewRegionJson)
}

// RegionMap is used to store mapping of country to region
type RegionMap map[string]uint64

// RegionJson defines a vindex that uses a lookup table.
// The table is expected to define the id column as unique. It's
// Unique and a Lookup.
type RegionJson struct {
	name        string
	regionMap   RegionMap
	regionBytes int
}

// NewRegionJson creates a RegionJson vindex.
// The supplied map requires all the fields of "RegionExperimental".
// Additionally, it requires a region_map argument representing the path to a json file
// containing a map of country to region.
func NewRegionJson(name string, m map[string]string) (Vindex, error) {
	rmPath := m["region_map"]
	rmap := make(map[string]uint64)
	data, err := ioutil.ReadFile(rmPath)
	if err != nil {
		return nil, err
	}
	log.Infof("Loaded Region map from: %s", rmPath)
	err = json.Unmarshal(data, &rmap)
	if err != nil {
		return nil, err
	}

	return &RegionJson{
		name:      name,
		regionMap: rmap,
	}, nil
}

// String returns the name of the vindex.
func (rv *RegionJson) String() string {
	return rv.name
}

// Cost returns the cost of this index as 1.
func (rv *RegionJson) Cost() int {
	return 1
}

// IsUnique returns true since the Vindex is unique.
func (rv *RegionJson) IsUnique() bool {
	return true
}

// Map satisfies MultiColumn.
func (rv *RegionJson) Map(vcursor VCursor, rowsColValues [][]sqltypes.Value) ([]key.Destination, error) {
	destinations := make([]key.Destination, 0, len(rowsColValues))
	for _, row := range rowsColValues {
		if len(row) != 2 {
			destinations = append(destinations, key.DestinationNone{})
			continue
		}
		// Compute hash.
		hn, err := sqltypes.ToUint64(row[0])
		if err != nil {
			destinations = append(destinations, key.DestinationNone{})
			continue
		}
		h := vhash(hn)

		rn, ok := rv.regionMap[row[1].ToString()]
		if !ok {
			destinations = append(destinations, key.DestinationNone{})
			continue
		}
		r := make([]byte, 2)
		binary.BigEndian.PutUint16(r, uint16(rn))

		// Concatenate and add to destinations.
		if rv.regionBytes == 1 {
			r = r[1:]
		}
		dest := append(r, h...)
		destinations = append(destinations, key.DestinationKeyspaceID(dest))
	}
	return destinations, nil
}

// Verify satisfies MultiColumn
func (rv *RegionJson) Verify(vcursor VCursor, rowsColValues [][]sqltypes.Value, ksids [][]byte) ([]bool, error) {
	result := make([]bool, len(rowsColValues))
	destinations, _ := rv.Map(vcursor, rowsColValues)
	for i, dest := range destinations {
		destksid, ok := dest.(key.DestinationKeyspaceID)
		if !ok {
			continue
		}
		result[i] = bytes.Equal([]byte(destksid), ksids[i])
	}
	return result, nil
}

// NeedVCursor satisfies the Vindex interface.
func (rv *RegionJson) NeedsVCursor() bool {
	return false
}
<|endoftext|>


<commit_before>package sendgrid

import (
	//"bytes"
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/smtp"
	"net/url"
	"path/filepath"
)

type SGClient struct {
	apiUser  string
	apiPwd   string
	apiUrl   string
	smtpUrl  string
	smtpPort string
	smtpAuth smtp.Auth
	// Client is the HTTP transport to use when making requests.
	// It will default to http.DefaultClient if nil.
	Client *http.Client
}

/*
apiUser - SG username
apiPwd - SG password
*/
func NewSendGridClient(apiUser, apiPwd string) SGClient {
	smtpUrl := "smtp.sendgrid.net"
	smtpPort := "587"
	apiUrl := "https://sendgrid.com/api/mail.send.json?"
	smtpAuth := smtp.PlainAuth("", apiUser, apiPwd, smtpUrl)
	return SGClient{
		apiUser:  apiUser,
		apiPwd:   apiPwd,
		apiUrl:   apiUrl,
		smtpUrl:  smtpUrl,
		smtpPort: smtpPort,
		smtpAuth: smtpAuth,
	}
}

/*
Send will try to use the WebAPI first. If it's a success then a nill will be returned.
Else, the SMTP API will be used as a fail over. If this happens regardless of what is the result
of the SMTP attempt, you will receive an array of errors.
If the length of the array is 1, then SMTP succeeded, else it also failed.
*/
func (sg *SGClient) Send(m Mail) []error {
	if apiError := sg.SendAPI(m); apiError != nil {
		var errors []error
		errors = append(errors, apiError)
		if smtpError := sg.SendSMTP(m); smtpError != nil {
			return append(errors, smtpError)
		} else {
			return errors
		}
	} else {
		return nil //sucess
	}
}

func (sg *SGClient) SendSMTP(m Mail) error {
	return smtp.SendMail(sg.smtpUrl+":"+sg.smtpPort, sg.smtpAuth, m.from, m.to, []byte(m.html))
}

func (sg *SGClient) SendAPI(m Mail) error {
	values := url.Values{}
	values.Set("api_user", sg.apiUser)
	values.Set("api_key", sg.apiPwd)
	values.Set("subject", m.subject)
	values.Set("html", m.html)
	values.Set("text", m.text)
	values.Set("from", m.from)
	for i := 0; i < len(m.to); i++ {
		values.Set("to[]", m.to[i])
	}
	for i := 0; i < len(m.bcc); i++ {
		values.Set("bcc[]", m.bcc[i])
	}
	for i := 0; i < len(m.toname); i++ {
		values.Set("toname[]", m.toname[i])
	}
	for k, v := range m.files {
		values.Set("files["+k+"]", v)
	}
	if sg.Client == nil {
		sg.Client = http.DefaultClient
	}
	r, e := sg.Client.PostForm(sg.apiUrl, values)
	defer r.Body.Close()
	if r.StatusCode == 200 && e == nil {
		return nil
	} else {
		body, _ := ioutil.ReadAll(r.Body)
		return fmt.Errorf("sendgrid.go: code:%d error:%v body:%s", r.StatusCode, e, body)
	}
}

type Mail struct {
	to       []string
	toname   []string
	subject  string
	html     string
	text     string
	from     string
	bcc      []string
	fromname string
	replyto  string
	date     string
	files    map[string]string
	//still missing some stuff
}

func NewMail() Mail {
	return Mail{}
}

/*
TODO: Validate email addressed with RegExp.
*/
func (m *Mail) AddTo(email string) {
	m.to = append(m.to, email)
}

func (m *Mail) AddToName(name string) {
	m.toname = append(m.toname, name)
}

func (m *Mail) AddSubject(s string) {
	m.subject = s
}

func (m *Mail) AddHTML(html string) {
	m.html = html
}

func (m *Mail) AddText(text string) {
	m.text = text
}

func (m *Mail) AddFrom(from string) {
	m.from = from
}

func (m *Mail) AddBCC(email string) {
	m.bcc = append(m.bcc, email)
}

func (m *Mail) AddFromName(name string) {
	m.fromname = name
}

func (m *Mail) AddReplyTo(reply string) {
	m.replyto = reply
}

func (m *Mail) AddDate(date string) {
	m.date = date
}

func (m *Mail) AddAttachment(filePath string) error {
	if m.files == nil {
		m.files = make(map[string]string)
	}
	buf, e := ioutil.ReadFile(filePath)
	if e != nil {
		return e
	}
	_, filename := filepath.Split(filePath)
	m.files[filename] = base64.StdEncoding.EncodeToString(buf)
	return nil
}
<commit_msg>Added header support<commit_after>package sendgrid

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/smtp"
	"net/url"
	"path/filepath"
)

type SGClient struct {
	apiUser  string
	apiPwd   string
	apiUrl   string
	smtpUrl  string
	smtpPort string
	smtpAuth smtp.Auth
	// Client is the HTTP transport to use when making requests.
	// It will default to http.DefaultClient if nil.
	Client *http.Client
}

/*
apiUser - SG username
apiPwd - SG password
*/
func NewSendGridClient(apiUser, apiPwd string) SGClient {
	smtpUrl := "smtp.sendgrid.net"
	smtpPort := "587"
	apiUrl := "https://sendgrid.com/api/mail.send.json?"
	smtpAuth := smtp.PlainAuth("", apiUser, apiPwd, smtpUrl)
	return SGClient{
		apiUser:  apiUser,
		apiPwd:   apiPwd,
		apiUrl:   apiUrl,
		smtpUrl:  smtpUrl,
		smtpPort: smtpPort,
		smtpAuth: smtpAuth,
	}
}

/*
Send will try to use the WebAPI first. If it's a success then a nill will be returned.
Else, the SMTP API will be used as a fail over. If this happens regardless of what is the result
of the SMTP attempt, you will receive an array of errors.
If the length of the array is 1, then SMTP succeeded, else it also failed.
*/
func (sg *SGClient) Send(m Mail) []error {
	if apiError := sg.SendAPI(m); apiError != nil {
		var errors []error
		errors = append(errors, apiError)
		if smtpError := sg.SendSMTP(m); smtpError != nil {
			return append(errors, smtpError)
		} else {
			return errors
		}
	} else {
		return nil //sucess
	}
}

func (sg *SGClient) SendSMTP(m Mail) error {
	return smtp.SendMail(sg.smtpUrl+":"+sg.smtpPort, sg.smtpAuth, m.from, m.to, []byte(m.html))
}

func (sg *SGClient) SendAPI(m Mail) error {
	values := url.Values{}
	values.Set("api_user", sg.apiUser)
	values.Set("api_key", sg.apiPwd)
	values.Set("subject", m.subject)
	values.Set("html", m.html)
	values.Set("text", m.text)
	values.Set("from", m.from)
	headers, e := json.Marshal(m.headers)
	if e != nil {
		return fmt.Errorf("sendgrid.go: Error parsing JSON headers")
	}
	values.Set("headers", string(headers[:]))
	for i := 0; i < len(m.to); i++ {
		values.Set("to[]", m.to[i])
	}
	for i := 0; i < len(m.bcc); i++ {
		values.Set("bcc[]", m.bcc[i])
	}
	for i := 0; i < len(m.toname); i++ {
		values.Set("toname[]", m.toname[i])
	}
	for k, v := range m.files {
		values.Set("files["+k+"]", v)
	}
	if sg.Client == nil {
		sg.Client = http.DefaultClient
	}
	fmt.Print(values)
	r, e := sg.Client.PostForm(sg.apiUrl, values)
	defer r.Body.Close()
	if r.StatusCode == 200 && e == nil {
		return nil
	} else {
		body, _ := ioutil.ReadAll(r.Body)
		return fmt.Errorf("sendgrid.go: code:%d error:%v body:%s", r.StatusCode, e, body)
	}
}

type Mail struct {
	to       []string
	toname   []string
	subject  string
	html     string
	text     string
	from     string
	bcc      []string
	fromname string
	replyto  string
	date     string
	files    map[string]string
	headers  map[string]string
	//still missing some stuff
}

func NewMail() Mail {
	return Mail{}
}

/*
TODO: Validate email addressed with RegExp.
*/
func (m *Mail) AddTo(email string) {
	m.to = append(m.to, email)
}

func (m *Mail) AddToName(name string) {
	m.toname = append(m.toname, name)
}

func (m *Mail) AddSubject(s string) {
	m.subject = s
}

func (m *Mail) AddHTML(html string) {
	m.html = html
}

func (m *Mail) AddText(text string) {
	m.text = text
}

func (m *Mail) AddFrom(from string) {
	m.from = from
}

func (m *Mail) AddBCC(email string) {
	m.bcc = append(m.bcc, email)
}

func (m *Mail) AddFromName(name string) {
	m.fromname = name
}

func (m *Mail) AddReplyTo(reply string) {
	m.replyto = reply
}

func (m *Mail) AddDate(date string) {
	m.date = date
}

func (m *Mail) AddHeader(header, value string) {
	if m.headers == nil {
		m.headers = make(map[string]string)
	}
	m.headers[header] = value
}

func (m *Mail) AddAttachment(filePath string) error {
	if m.files == nil {
		m.files = make(map[string]string)
	}
	buf, e := ioutil.ReadFile(filePath)
	if e != nil {
		return e
	}
	_, filename := filepath.Split(filePath)
	m.files[filename] = base64.StdEncoding.EncodeToString(buf)
	return nil
}
<|endoftext|>