<commit_before>// Copyright (c) 2015 Mute Communications Ltd.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package msg

import (
	"bytes"
	"crypto/sha512"
	"io"

	"github.com/mutecomm/mute/cipher"
	"github.com/mutecomm/mute/encode/base64"
	"github.com/mutecomm/mute/msg/session"
	"github.com/mutecomm/mute/util/bzero"
	"golang.org/x/crypto/hkdf"
)

// checkKeys checks that the keys k1, k2, k3, and k4 are pairwise different to
// prevent possible reflection attacks.
func checkKeys(k1, k2, k3, k4 *[32]byte) error {
	if bytes.Equal(k1[:], k2[:]) {
		return ErrReflection
	}
	if bytes.Equal(k1[:], k3[:]) {
		return ErrReflection
	}
	if bytes.Equal(k1[:], k4[:]) {
		return ErrReflection
	}
	if bytes.Equal(k2[:], k3[:]) {
		return ErrReflection
	}
	if bytes.Equal(k2[:], k4[:]) {
		return ErrReflection
	}
	if bytes.Equal(k3[:], k4[:]) {
		return ErrReflection
	}
	return nil
}

// deriveRootKey derives the next root key from t1, t2, t3, and the
// previousRootKeyHash (if it exists).
func deriveRootKey(
	t1, t2, t3 *[32]byte,
	previousRootKeyHash *[64]byte,
) (*[24]byte, error) {
	master := make([]byte, 32+32+32+64)
	copy(master[:], t1[:])
	copy(master[32:], t2[:])
	copy(master[64:], t3[:])
	if previousRootKeyHash != nil {
		copy(master[96:], previousRootKeyHash[:])
	}

	hkdf := hkdf.New(sha512.New, master, nil, nil)

	// derive root key
	// TODO: size correct? Shouldn't it be 64 bytes?
	var rootKey [24]byte
	if _, err := io.ReadFull(hkdf, rootKey[:]); err != nil {
		return nil, err
	}

	return &rootKey, nil
}

// generateMessageKeys generates the next numOfKeys many session keys from
// from rootKey for given senderIdentity and recipientIdentity.
// If recipientKeys is true the generated sender and reciever keys are stored in
// reverse order.
// It uses senderSessionPub and recipientPub in the process and calls
// keyStore.StoresSession and keyStore.SetSessionState to store the result.
func generateMessageKeys(
	senderIdentity, recipientIdentity string,
	rootKey []byte, // TODO: change to *[64]byte?
	recipientKeys bool,
	senderSessionPub, recipientPub *[32]byte,
	numOfKeys uint64,
	keyStore session.Store,
) error {
	var (
		identities string
		send       []string
		recv       []string
	)

	// identity_fix = HASH(SORT(SenderNym, RecipientNym))
	if senderIdentity < recipientIdentity {
		identities = senderIdentity + recipientIdentity
	} else {
		identities = recipientIdentity + senderIdentity
	}
	identityFix := cipher.SHA512([]byte(identities))
	recipientPubHash := cipher.SHA512(recipientPub[:])
	senderSessionPubHash := cipher.SHA512(senderSessionPub[:])

	chainKey := rootKey
	for i := uint64(0); i < numOfKeys; i++ {
		// messagekey_send[i] = HMAC_HASH(chainkey, "MESSAGE" | HASH(RecipientPub) | identity_fix)
		buffer := append([]byte("MESSAGE"), recipientPubHash...)
		buffer = append(buffer, identityFix...)
		send = append(send, base64.Encode(cipher.HMAC(chainKey, buffer)))

		// messagekey_recv[i] = HMAC_HASH(chainkey, "MESSAGE" | HASH(SenderSessionPub) | identity_fix)
		buffer = append([]byte("MESSAGE"), senderSessionPubHash...)
		buffer = append(buffer, identityFix...)
		recv = append(recv, base64.Encode(cipher.HMAC(chainKey, buffer)))

		// chainkey = HMAC_HASH(chainkey, "CHAIN" )
		chainKey = cipher.HMAC(chainKey, []byte("CHAIN"))
	}

	// calculate root key hash
	rootKeyHash := base64.Encode(cipher.SHA512(rootKey[:]))
	bzero.Bytes(rootKey[:])

	// reverse key material, if necessary
	if recipientKeys {
		senderIdentity, recipientIdentity = recipientIdentity, senderIdentity
		send, recv = recv, send
	}

	// store session
	var pubHash string
	if recipientKeys {
		pubHash = base64.Encode(cipher.SHA512(recipientPub[:]))
	} else {
		pubHash = base64.Encode(cipher.SHA512(senderSessionPub[:]))
	}
	err := keyStore.StoreSession(senderIdentity, recipientIdentity, pubHash,
		rootKeyHash, base64.Encode(chainKey), send, recv)
	if err != nil {
		return err
	}

	return nil
}

// deriveSymmetricKeys derives the symmetric cryptoKey and hmacKey from the
// given messageKey.
func deriveSymmetricKeys(messageKey *[64]byte) (
	cryptoKey, hmacKey []byte,
	err error,
) {
	// TODO: set optional salt and info?
	hkdf := hkdf.New(sha512.New, messageKey[:], nil, nil)

	// derive crypto key for AES-256
	cryptoKey = make([]byte, 32)
	if _, err := io.ReadFull(hkdf, cryptoKey); err != nil {
		return nil, nil, err
	}

	// derive HMAC key for SHA-512 HMAC (TODO: correct size?)
	hmacKey = make([]byte, 64)
	if _, err := io.ReadFull(hkdf, hmacKey); err != nil {
		return nil, nil, err
	}

	return
}
<commit_msg>review msg/kesy.go<commit_after>// Copyright (c) 2015 Mute Communications Ltd.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package msg

import (
	"bytes"
	"crypto/sha512"
	"io"

	"github.com/mutecomm/mute/cipher"
	"github.com/mutecomm/mute/encode/base64"
	"github.com/mutecomm/mute/msg/session"
	"github.com/mutecomm/mute/util/bzero"
	"golang.org/x/crypto/hkdf"
)

// checkKeys checks that the keys k1, k2, k3, and k4 are pairwise different to
// prevent possible reflection attacks.
func checkKeys(k1, k2, k3, k4 *[32]byte) error {
	if bytes.Equal(k1[:], k2[:]) {
		return ErrReflection
	}
	if bytes.Equal(k1[:], k3[:]) {
		return ErrReflection
	}
	if bytes.Equal(k1[:], k4[:]) {
		return ErrReflection
	}
	if bytes.Equal(k2[:], k3[:]) {
		return ErrReflection
	}
	if bytes.Equal(k2[:], k4[:]) {
		return ErrReflection
	}
	if bytes.Equal(k3[:], k4[:]) {
		return ErrReflection
	}
	return nil
}

// deriveRootKey derives the next root key from t1, t2, t3, and the
// previousRootKeyHash (if it exists).
func deriveRootKey(
	t1, t2, t3 *[32]byte,
	previousRootKeyHash *[64]byte,
) (*[24]byte, error) {
	master := make([]byte, 32+32+32+64)
	copy(master[:], t1[:])
	copy(master[32:], t2[:])
	copy(master[64:], t3[:])
	if previousRootKeyHash != nil {
		copy(master[96:], previousRootKeyHash[:])
	}

	hkdf := hkdf.New(sha512.New, master, nil, nil)

	// derive root key
	// TODO: size correct? Shouldn't it be 64 bytes? -> 32 byte (we operate on 256bit security margins)
	var rootKey [24]byte
	if _, err := io.ReadFull(hkdf, rootKey[:]); err != nil {
		return nil, err
	}

	return &rootKey, nil
}

// generateMessageKeys generates the next numOfKeys many session keys from
// from rootKey for given senderIdentity and recipientIdentity.
// If recipientKeys is true the generated sender and reciever keys are stored in
// reverse order.
// It uses senderSessionPub and recipientPub in the process and calls
// keyStore.StoresSession and keyStore.SetSessionState to store the result.
func generateMessageKeys(
	senderIdentity, recipientIdentity string,
	rootKey []byte, // TODO: change to *[64]byte? -> 32 byte (we operate on 256 bit security margins)
	recipientKeys bool,
	senderSessionPub, recipientPub *[32]byte,
	numOfKeys uint64,
	keyStore session.Store,
) error {
	var (
		identities string
		send       []string
		recv       []string
	)

	// identity_fix = HASH(SORT(SenderNym, RecipientNym))
	if senderIdentity < recipientIdentity {
		identities = senderIdentity + recipientIdentity
	} else {
		identities = recipientIdentity + senderIdentity
	}
	identityFix := cipher.SHA512([]byte(identities))
	recipientPubHash := cipher.SHA512(recipientPub[:])
	senderSessionPubHash := cipher.SHA512(senderSessionPub[:])

	chainKey := rootKey
	for i := uint64(0); i < numOfKeys; i++ {
		// messagekey_send[i] = HMAC_HASH(chainkey, "MESSAGE" | HASH(RecipientPub) | identity_fix)
		buffer := append([]byte("MESSAGE"), recipientPubHash...)
		buffer = append(buffer, identityFix...)
		send = append(send, base64.Encode(cipher.HMAC(chainKey, buffer)))

		// messagekey_recv[i] = HMAC_HASH(chainkey, "MESSAGE" | HASH(SenderSessionPub) | identity_fix)
		buffer = append([]byte("MESSAGE"), senderSessionPubHash...)
		buffer = append(buffer, identityFix...)
		recv = append(recv, base64.Encode(cipher.HMAC(chainKey, buffer)))

		// chainkey = HMAC_HASH(chainkey, "CHAIN" )
		chainKey = cipher.HMAC(chainKey, []byte("CHAIN"))
	}

	// calculate root key hash
	rootKeyHash := base64.Encode(cipher.SHA512(rootKey[:]))
	bzero.Bytes(rootKey[:])

	// reverse key material, if necessary
	if recipientKeys {
		senderIdentity, recipientIdentity = recipientIdentity, senderIdentity
		send, recv = recv, send
	}

	// store session
	var pubHash string
	if recipientKeys {
		pubHash = base64.Encode(cipher.SHA512(recipientPub[:]))
	} else {
		pubHash = base64.Encode(cipher.SHA512(senderSessionPub[:]))
	}
	err := keyStore.StoreSession(senderIdentity, recipientIdentity, pubHash,
		rootKeyHash, base64.Encode(chainKey), send, recv)
	if err != nil {
		return err
	}

	return nil
}

// deriveSymmetricKeys derives the symmetric cryptoKey and hmacKey from the
// given messageKey.
func deriveSymmetricKeys(messageKey *[64]byte) (
	cryptoKey, hmacKey []byte,
	err error,
) {
	// TODO: set optional salt and info? -> no
	hkdf := hkdf.New(sha512.New, messageKey[:], nil, nil)

	// derive crypto key for AES-256
	cryptoKey = make([]byte, 32)
	if _, err := io.ReadFull(hkdf, cryptoKey); err != nil {
		return nil, nil, err
	}

	// derive HMAC key for SHA-512 HMAC (TODO: correct size? -> only 32 byte)
	hmacKey = make([]byte, 64)
	if _, err := io.ReadFull(hkdf, hmacKey); err != nil {
		return nil, nil, err
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"log"
	"net/http"
	"sort"
	"strings"
)

// LoadWordList loads in a list of known-good words
const scowlUS60 = "http://app.aspell.net/create?max_size=60&spelling=US&max_variant=1&diacritic=strip&special=hacker&download=wordlist&encoding=utf-8&format=inline"

func loadWordList(target string) map[string]bool {
	resp, err := http.Get(target)
	if err != nil {
		log.Fatalf("Unable to download: %s", err)
	}
	defer resp.Body.Close()
	out := make(map[string]bool)
	intro := true
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		line := scanner.Text()
		if intro {
			if line == "---" {
				intro = false
			}
			continue
		}
		out[strings.ToLower(line)] = true
	}
	if err := scanner.Err(); err != nil {
		log.Fatalf("Error reading: %s", err)
	}
	return out
}

func main() {
	wordmap := loadWordList(scowlUS60)
	// do any fixups here

	words := make([]string, 0, len(wordmap))
	for k, _ := range wordmap {
		words = append(words, k)
	}
	log.Printf("Got %d words", len(words))

	sort.Strings(words)
	for _, word := range words {
		fmt.Println(word)
	}
}
<commit_msg>productionalize<commit_after>package main

import (
	"bufio"
	"flag"
	"log"
	"net/http"
	"os"
	"sort"
	"strings"
)

// LoadWordList loads in a list of known-good words
const scowlUS60 = "http://app.aspell.net/create?max_size=60&spelling=US&max_variant=1&diacritic=strip&special=hacker&download=wordlist&encoding=utf-8&format=inline"

func loadWordList(target string) map[string]bool {
	resp, err := http.Get(target)
	if err != nil {
		log.Fatalf("Unable to download: %s", err)
	}
	defer resp.Body.Close()
	out := make(map[string]bool)
	intro := true
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		line := scanner.Text()
		if intro {
			if line == "---" {
				intro = false
			}
			continue
		}

		// we dont use contractions
		if strings.Contains(line, "'") {
			continue
		}

		out[strings.ToLower(line)] = true
	}
	if err := scanner.Err(); err != nil {
		log.Fatalf("Error reading: %s", err)
	}
	return out
}

func main() {
	outfile := flag.String("out", "words.txt", "Name of output file")
	flag.Parse()

	wordmap := loadWordList(scowlUS60)
	// do any fixups here

	words := make([]string, 0, len(wordmap))
	for k, _ := range wordmap {
		words = append(words, k)
	}
	log.Printf("Got %d words", len(words))

	sort.Strings(words)

	fo, err := os.Create(*outfile)
	if err != nil {
		log.Fatalf("unable to create %s: %s", *outfile, err)
	}

	buf := bufio.NewWriter(fo)
	for _, word := range words {
		buf.WriteString(word + "\n")
	}
	buf.Flush()
	fo.Close()
}
<|endoftext|>


<commit_before>package sls

import (
	"encoding/json"
	"fmt"
	"github.com/denverdino/aliyungo/common"
	"github.com/golang/protobuf/proto"
	"net/http"
	"time"
	"strconv"
)

type Client struct {
	accessKeyId     string //Access Key Id
	accessKeySecret string //Access Key Secret
	debug           bool
	httpClient      *http.Client
	version         string
	internal        bool
	region          common.Region
	endpoint        string
}

type Project struct {
	client      *Client
	Name        string `json:"projectName,omitempty"`
	Description string `json:"description,omitempty"`
}

type LogContent struct {
	Key   string
	Value string
}

type LogItem struct {
	Time     time.Time
	Contents []*LogContent
}

type LogGroupItem struct {
	Logs   []*LogItem
	Topic  string
	Source string
}

type PutLogsRequest struct {
	Project  string
	LogStore string
	LogItems LogGroup
	HashKey  string
}

const (
	SLSDefaultEndpoint = "sls.aliyuncs.com"
	SLSAPIVersion      = "0.6.0"
	METHOD_GET         = "GET"
	METHOD_POST        = "POST"
	METHOD_PUT         = "PUT"
	METHOD_DELETE      = "DELETE"
)

// NewClient creates a new instance of ECS client
func NewClient(region common.Region, internal bool, accessKeyId, accessKeySecret string) *Client {
	return &Client{
		accessKeyId:     accessKeyId,
		accessKeySecret: accessKeySecret,
		internal:        internal,
		region:          region,
		version:         SLSAPIVersion,
		endpoint:        SLSDefaultEndpoint,
		httpClient:      &http.Client{},
	}
}

func (client *Client) Project(name string) (*Project, error) {

	//	newClient := client.forProject(name)
	//
	//	req := &request{
	//		method: METHOD_GET,
	//		path: "/",
	//	}
	//
	//	project := &Project{}
	//
	//	if err := newClient.requestWithJsonResponse(req, project); err != nil {
	//		return nil, err
	//	}
	//	project.client = newClient
	//	return project, nil
	return &Project{
		Name:   name,
		client: client.forProject(name),
	}, nil
}

func (client *Client) forProject(name string) *Client {
	newclient := *client

	region := string(client.region)
	if client.internal {
		region = fmt.Sprintf("%s-intranet", region)
	}
	newclient.endpoint = fmt.Sprintf("%s.%s.%s", name, region, SLSDefaultEndpoint)
	return &newclient
}

func (client *Client) DeleteProject(name string) error {
	req := &request{
		method: METHOD_DELETE,
		path:   "/",
	}

	newClient := client.forProject(name)
	return newClient.requestWithClose(req)
}

func (client *Client) CreateProject(name string, description string) error {
	project := &Project{
		Name:        name,
		Description: description,
	}
	data, err := json.Marshal(project)
	if err != nil {
		return err
	}

	req := &request{
		method:      METHOD_POST,
		path:        "/",
		payload:     data,
		contentType: "application/json",
	}

	newClient := client.forProject(name)
	return newClient.requestWithClose(req)
}

//
//func marshal() ([]byte, error) {
//
//	logGroups := []*LogGroup{}
//	tmp := []*LogGroupItem
//	for _, logGroupItem := range tmp {
//
//		logs := []*Log{}
//		for _, logItem := range logGroupItem.Logs {
//			contents := []*Log_Content{}
//			for key, value := range logItem.Content {
//				contents = append(contents, &Log_Content{
//					Key: proto.String(key),
//					Value: proto.String(value),
//				})
//			}
//
//			logs = append(logs, &Log{
//				Time: proto.Uint32(uint32(LogItem.Time.Unix())),
//				Contents: contents,
//			})
//		}
//
//		logGroup := &LogGroup{
//			Topic: proto.String(LogGroupItem.Topic),
//			Source: proto.String(LogGroupItem.Source),
//			Logs: logs,
//		}
//		logGroups = append(logGroups, logGroup)
//	}
//
//	return proto.Marshal(&LogGroupList{
//		LogGroupList: logGroups,
//	})
//}

func (client *Client) PutLogs(putLogRequest *PutLogsRequest) error {
	if putLogRequest == nil {
		return nil
	}

	data, err := proto.Marshal(&putLogRequest.LogItems)
	if err != nil {
		return err
	}

	req := &request{
		method:      METHOD_POST,
		path:        "/logstores/" + putLogRequest.LogStore + "/shards/lb",
		payload:     data,
		contentType: "application/x-protobuf",
		headers: map[string]string{ "x-log-bodyrawsize" : strconv.Itoa( len(data) ), },
	}

	newClient := client.forProject(putLogRequest.Project )
	return newClient.requestWithClose(req)

}
<commit_msg>clean code<commit_after>package sls

import (
	"encoding/json"
	"fmt"
	"github.com/denverdino/aliyungo/common"
	"github.com/golang/protobuf/proto"
	"net/http"
	//"time"
	"strconv"
)

type Client struct {
	accessKeyId     string //Access Key Id
	accessKeySecret string //Access Key Secret
	debug           bool
	httpClient      *http.Client
	version         string
	internal        bool
	region          common.Region
	endpoint        string
}

type Project struct {
	client      *Client
	Name        string `json:"projectName,omitempty"`
	Description string `json:"description,omitempty"`
}

//type LogContent struct {
//	Key   string
//	Value string
//}
//
//type LogItem struct {
//	Time     time.Time
//	Contents []*LogContent
//}
//
//type LogGroupItem struct {
//	Logs   []*LogItem
//	Topic  string
//	Source string
//}

type PutLogsRequest struct {
	Project  string
	LogStore string
	LogItems LogGroup
	HashKey  string
}

const (
	SLSDefaultEndpoint = "sls.aliyuncs.com"
	SLSAPIVersion      = "0.6.0"
	METHOD_GET         = "GET"
	METHOD_POST        = "POST"
	METHOD_PUT         = "PUT"
	METHOD_DELETE      = "DELETE"
)

// NewClient creates a new instance of ECS client
func NewClient(region common.Region, internal bool, accessKeyId, accessKeySecret string) *Client {
	return &Client{
		accessKeyId:     accessKeyId,
		accessKeySecret: accessKeySecret,
		internal:        internal,
		region:          region,
		version:         SLSAPIVersion,
		endpoint:        SLSDefaultEndpoint,
		httpClient:      &http.Client{},
	}
}

func (client *Client) Project(name string) (*Project, error) {

	//	newClient := client.forProject(name)
	//
	//	req := &request{
	//		method: METHOD_GET,
	//		path: "/",
	//	}
	//
	//	project := &Project{}
	//
	//	if err := newClient.requestWithJsonResponse(req, project); err != nil {
	//		return nil, err
	//	}
	//	project.client = newClient
	//	return project, nil
	return &Project{
		Name:   name,
		client: client.forProject(name),
	}, nil
}

func (client *Client) forProject(name string) *Client {
	newclient := *client

	region := string(client.region)
	if client.internal {
		region = fmt.Sprintf("%s-intranet", region)
	}
	newclient.endpoint = fmt.Sprintf("%s.%s.%s", name, region, SLSDefaultEndpoint)
	return &newclient
}

func (client *Client) DeleteProject(name string) error {
	req := &request{
		method: METHOD_DELETE,
		path:   "/",
	}

	newClient := client.forProject(name)
	return newClient.requestWithClose(req)
}

func (client *Client) CreateProject(name string, description string) error {
	project := &Project{
		Name:        name,
		Description: description,
	}
	data, err := json.Marshal(project)
	if err != nil {
		return err
	}

	req := &request{
		method:      METHOD_POST,
		path:        "/",
		payload:     data,
		contentType: "application/json",
	}

	newClient := client.forProject(name)
	return newClient.requestWithClose(req)
}

//
//func marshal() ([]byte, error) {
//
//	logGroups := []*LogGroup{}
//	tmp := []*LogGroupItem
//	for _, logGroupItem := range tmp {
//
//		logs := []*Log{}
//		for _, logItem := range logGroupItem.Logs {
//			contents := []*Log_Content{}
//			for key, value := range logItem.Content {
//				contents = append(contents, &Log_Content{
//					Key: proto.String(key),
//					Value: proto.String(value),
//				})
//			}
//
//			logs = append(logs, &Log{
//				Time: proto.Uint32(uint32(LogItem.Time.Unix())),
//				Contents: contents,
//			})
//		}
//
//		logGroup := &LogGroup{
//			Topic: proto.String(LogGroupItem.Topic),
//			Source: proto.String(LogGroupItem.Source),
//			Logs: logs,
//		}
//		logGroups = append(logGroups, logGroup)
//	}
//
//	return proto.Marshal(&LogGroupList{
//		LogGroupList: logGroups,
//	})
//}

func (client *Client) PutLogs(putLogRequest *PutLogsRequest) error {
	if putLogRequest == nil {
		return nil
	}

	data, err := proto.Marshal(&putLogRequest.LogItems)
	if err != nil {
		return err
	}

	req := &request{
		method:      METHOD_POST,
		path:        "/logstores/" + putLogRequest.LogStore + "/shards/lb",
		payload:     data,
		contentType: "application/x-protobuf",
		headers: map[string]string{ "x-log-bodyrawsize" : strconv.Itoa( len(data) ), },
	}

	newClient := client.forProject(putLogRequest.Project )
	return newClient.requestWithClose(req)

}
<|endoftext|>


<commit_before>// Copyright 2013-2014 Adam Presley. All rights reserved
// Use of this source code is governed by the MIT license
// that can be found in the LICENSE file.

package smtp

import (
	"database/sql"
	"fmt"
	"log"

	_ "github.com/go-sql-driver/mysql"
)

func ConnectMySQL(host string, port string, database string, userName string, password string) (*sql.DB, error) {
	/*
	 * Create the connection
	 */
	log.Println("Connecting to MySQL database")

	db, err := sql.Open("mysql", fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?autocommit=true", userName, password, host, port, database))
	if err != nil {
		return nil, err
	}

	return db, nil
}

func CreateMySQLDatabase(db *sql.DB) error {
	log.Println("Creating tables...")

	var err error

	sql := `
		CREATE TABLE IF NOT EXISTS mailitem (
			id INT NOT NULL PRIMARY KEY AUTO_INCREMENT,
			dateSent DATETIME,
			fromAddress VARCHAR(255),
			toAddressList TEXT,
			subject VARCHAR(512),
			xmailer VARCHAR(50),
			body TEXT,
			contentType VARCHAR(50),
			boundary VARCHAR(50)
		);
	`

	_, err = db.Exec(sql)
	if err != nil {
		return err
	}

	sql = `
		CREATE TABLE IF NOT EXISTS attachment (
			id INT NOT NULL PRIMARY KEY AUTO_INCREMENT,
			mailItemId INT,
			fileName VARCHAR(255),
			contentType VARCHAR(50),
			content TEXT
		);
	`

	_, err = db.Exec(sql)
	if err != nil {
		return err
	}

	log.Println("Created tables successfully.")
	return nil
}
<commit_msg>Setting max idle and open connections for MySQL based on environment variable "max_connections". Defaults to MySQL 5 value of 151. refs #2<commit_after>// Copyright 2013-2014 Adam Presley. All rights reserved
// Use of this source code is governed by the MIT license
// that can be found in the LICENSE file.

package smtp

import (
	"database/sql"
	"fmt"
	"log"
	"os"
	"strconv"

	_ "github.com/go-sql-driver/mysql"
)

func ConnectMySQL(host string, port string, database string, userName string, password string) (*sql.DB, error) {
	/*
	 * Create the connection
	 */
	log.Println("Connecting to MySQL database")

	db, err := sql.Open("mysql", fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?autocommit=true", userName, password, host, port, database))
	if err != nil {
		return nil, err
	}

	temp := os.Getenv("max_connections")
	if temp == "" {
		temp = "151"
	}

	maxConnections, _ := strconv.Atoi(temp)

	db.SetMaxIdleConns(maxConnections)
	db.SetMaxOpenConns(maxConnections)

	return db, nil
}

func CreateMySQLDatabase(db *sql.DB) error {
	log.Println("Creating tables...")

	var err error

	sql := `
		CREATE TABLE IF NOT EXISTS mailitem (
			id INT NOT NULL PRIMARY KEY AUTO_INCREMENT,
			dateSent DATETIME,
			fromAddress VARCHAR(255),
			toAddressList TEXT,
			subject VARCHAR(512),
			xmailer VARCHAR(50),
			body TEXT,
			contentType VARCHAR(50),
			boundary VARCHAR(50)
		);
	`

	_, err = db.Exec(sql)
	if err != nil {
		return err
	}

	sql = `
		CREATE TABLE IF NOT EXISTS attachment (
			id INT NOT NULL PRIMARY KEY AUTO_INCREMENT,
			mailItemId INT,
			fileName VARCHAR(255),
			contentType VARCHAR(50),
			content TEXT
		);
	`

	_, err = db.Exec(sql)
	if err != nil {
		return err
	}

	log.Println("Created tables successfully.")
	return nil
}
<|endoftext|>


<commit_before>package services

import (
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/google/go-github/github"
	"github.com/matrix-org/go-neb/database"
	"github.com/matrix-org/go-neb/matrix"
	"github.com/matrix-org/go-neb/plugin"
	"github.com/matrix-org/go-neb/realms/github"
	"github.com/matrix-org/go-neb/services/github/client"
	"github.com/matrix-org/go-neb/services/github/webhook"
	"github.com/matrix-org/go-neb/types"
	"net/http"
	"regexp"
	"strconv"
	"strings"
)

// Matches alphanumeric then a /, then more alphanumeric then a #, then a number.
// E.g. owner/repo#11 (issue/PR numbers) - Captured groups for owner/repo/number
var ownerRepoIssueRegex = regexp.MustCompile("([A-z0-9-_]+)/([A-z0-9-_]+)#([0-9]+)")

type githubService struct {
	id                 string
	webhookEndpointURL string
	BotUserID          string
	ClientUserID       string
	RealmID            string
	SecretToken        string
	Rooms              map[string]struct { // room_id => {}
		Repos map[string]struct { // owner/repo => { events: ["push","issue","pull_request"] }
			Events []string
		}
	}
}

func (s *githubService) ServiceUserID() string { return s.BotUserID }
func (s *githubService) ServiceID() string     { return s.id }
func (s *githubService) ServiceType() string   { return "github" }
func (s *githubService) RoomIDs() []string {
	var keys []string
	for k := range s.Rooms {
		keys = append(keys, k)
	}
	return keys
}

func (s *githubService) cmdGithubCreate(roomID, userID string, args []string) (interface{}, error) {
	cli := s.githubClientFor(userID, false)
	if cli == nil {
		r, err := database.GetServiceDB().LoadAuthRealm(s.RealmID)
		if err != nil {
			return nil, err
		}
		ghRealm, ok := r.(*realms.GithubRealm)
		if !ok {
			return nil, fmt.Errorf("Failed to cast realm %s into a GithubRealm", s.RealmID)
		}
		return matrix.StarterLinkMessage{
			Body: "You need to OAuth with Github before you can create issues.",
			Link: ghRealm.StarterLink,
		}, nil
	}

	if len(args) < 2 {
		return &matrix.TextMessage{"m.notice",
			`Usage: !github create owner/repo "issue title" "description"`}, nil
	}

	var (
		ownerRepo string
		title     *string
		desc      *string
	)
	ownerRepo = args[0]
	o := strings.Split(ownerRepo, "/")
	if len(o) != 2 {
		return &matrix.TextMessage{"m.notice",
			`Usage: !github create owner/repo "issue title" "description"`}, nil
	}

	if len(args) == 2 {
		title = &args[1]
	} else if len(args) == 3 {
		title = &args[1]
		desc = &args[2]
	} else { // > 3 args is probably a title without quote marks
		joinedTitle := strings.Join(args[1:], " ")
		title = &joinedTitle
	}

	issue, res, err := cli.Issues.Create(o[0], o[1], &github.IssueRequest{
		Title: title,
		Body:  desc,
	})
	if err != nil {
		log.WithField("err", err).Print("Failed to create issue")
		return nil, fmt.Errorf("Failed to create issue. HTTP %d", res.StatusCode)
	}

	return matrix.TextMessage{"m.notice", fmt.Sprintf("Created issue: %s", *issue.HTMLURL)}, nil
}

func (s *githubService) expandIssue(roomID, userID string, matchingGroups []string) interface{} {
	// matchingGroups => ["foo/bar#11", "foo", "bar", "11"]
	if len(matchingGroups) != 4 {
		log.WithField("groups", matchingGroups).Print("Unexpected number of groups")
		return nil
	}
	num, err := strconv.Atoi(matchingGroups[3])
	if err != nil {
		log.WithField("issue_number", matchingGroups[3]).Print("Bad issue number")
		return nil
	}
	owner := matchingGroups[1]
	repo := matchingGroups[2]

	cli := s.githubClientFor(userID, true)

	i, _, err := cli.Issues.Get(owner, repo, num)
	if err != nil {
		log.WithError(err).WithFields(log.Fields{
			"owner":  owner,
			"repo":   repo,
			"number": num,
		}).Print("Failed to fetch issue")
		return nil
	}

	return &matrix.TextMessage{
		"m.notice",
		fmt.Sprintf("%s : %s", *i.HTMLURL, *i.Title),
	}
}

func (s *githubService) Plugin(roomID string) plugin.Plugin {
	return plugin.Plugin{
		Commands: []plugin.Command{
			plugin.Command{
				Path: []string{"github", "create"},
				Command: func(roomID, userID string, args []string) (interface{}, error) {
					return s.cmdGithubCreate(roomID, userID, args)
				},
			},
		},
		Expansions: []plugin.Expansion{
			plugin.Expansion{
				Regexp: ownerRepoIssueRegex,
				Expand: func(roomID, userID string, matchingGroups []string) interface{} {
					return s.expandIssue(roomID, userID, matchingGroups)
				},
			},
		},
	}
}
func (s *githubService) OnReceiveWebhook(w http.ResponseWriter, req *http.Request, cli *matrix.Client) {
	evType, repo, msg, err := webhook.OnReceiveRequest(req, s.SecretToken)
	if err != nil {
		w.WriteHeader(err.Code)
		return
	}

	for roomID, roomConfig := range s.Rooms {
		for ownerRepo, repoConfig := range roomConfig.Repos {
			if !strings.EqualFold(*repo.FullName, ownerRepo) {
				continue
			}

			notifyRoom := false
			for _, notifyType := range repoConfig.Events {
				if evType == notifyType {
					notifyRoom = true
					break
				}
			}
			if notifyRoom {
				log.WithFields(log.Fields{
					"type":    evType,
					"msg":     msg,
					"repo":    repo,
					"room_id": roomID,
				}).Print("Sending notification to room")
				_, e := cli.SendMessageEvent(roomID, "m.room.message", msg)
				if e != nil {
					log.WithError(e).WithField("room_id", roomID).Print(
						"Failed to send notification to room.")
				}
			}
		}
	}
	w.WriteHeader(200)
}
func (s *githubService) Register() error {
	if s.RealmID == "" || s.ClientUserID == "" || s.BotUserID == "" {
		return fmt.Errorf("RealmID, BotUserID and ClientUserID are required")
	}
	// check realm exists
	realm, err := database.GetServiceDB().LoadAuthRealm(s.RealmID)
	if err != nil {
		return err
	}
	// make sure the realm is of the type we expect
	if realm.Type() != "github" {
		return fmt.Errorf("Realm is of type '%s', not 'github'", realm.Type())
	}

	// In order to register the GH service, you must have authed with GH.
	cli := s.githubClientFor(s.ClientUserID, false)
	if cli == nil {
		return fmt.Errorf(
			"User %s does not have a Github auth session with realm %s.", s.ClientUserID, realm.ID())
	}

	return nil
}

func (s *githubService) PostRegister(oldService types.Service) {
	cli := s.githubClientFor(s.ClientUserID, false)
	if cli == nil {
		log.Errorf("PostRegister: %s does not have a github session", s.ClientUserID)
		return
	}
	if oldService != nil {
		old, ok := oldService.(*githubService)
		if !ok {
			log.Error("PostRegister: Provided old service is not of type GithubService")
			return
		}

		// TODO: We should be adding webhooks in Register() then removing old hooks in PostRegister()
		//
		// By doing both operations in PostRegister(), if some of the requests fail we can end up in
		// an inconsistent state. It is a lot simpler and easy to reason about this way though, so
		// for now it will do.

		// remove any existing webhooks this service created on the user's behalf
		modifyWebhooks(old, cli, true)
	}

	// make new webhooks according to service config
	modifyWebhooks(s, cli, false)
}

func modifyWebhooks(s *githubService, cli *github.Client, removeHooks bool) {
	ownerRepoSet := make(map[string]bool)
	for _, roomCfg := range s.Rooms {
		for ownerRepo := range roomCfg.Repos {
			// sanity check that it looks like 'owner/repo' as we'll split on / later
			if strings.Count(ownerRepo, "/") != 1 {
				log.WithField("owner_repo", ownerRepo).Print("Bad owner/repo value.")
				continue
			}
			ownerRepoSet[ownerRepo] = true
		}
	}

	for ownerRepo := range ownerRepoSet {
		o := strings.Split(ownerRepo, "/")
		owner := o[0]
		repo := o[1]
		logger := log.WithFields(log.Fields{
			"owner": owner,
			"repo":  repo,
		})
		if removeHooks {
			removeHook(logger, cli, owner, repo, s.webhookEndpointURL)
		} else {
			// make a hook for all GH events since we'll filter it when we receive webhook requests
			name := "web" // https://developer.github.com/v3/repos/hooks/#create-a-hook
			cfg := map[string]interface{}{
				"content_type": "json",
				"url":          s.webhookEndpointURL,
			}
			if s.SecretToken != "" {
				cfg["secret"] = s.SecretToken
			}
			events := []string{"push", "pull_request", "issues", "issue_comment", "pull_request_review_comment"}
			_, _, err := cli.Repositories.CreateHook(owner, repo, &github.Hook{
				Name:   &name,
				Config: cfg,
				Events: events,
			})
			if err != nil {
				logger.WithError(err).Print("Failed to create webhook")
				// continue as others may succeed
			}
		}
	}
}

func (s *githubService) githubClientFor(userID string, allowUnauth bool) *github.Client {
	token, err := getTokenForUser(s.RealmID, userID)
	if err != nil {
		log.WithFields(log.Fields{
			log.ErrorKey: err,
			"user_id":    userID,
			"realm_id":   s.RealmID,
		}).Print("Failed to get token for user")
	}
	if token != "" {
		return client.New(token)
	} else if allowUnauth {
		return client.New("")
	} else {
		return nil
	}
}

func getTokenForUser(realmID, userID string) (string, error) {
	realm, err := database.GetServiceDB().LoadAuthRealm(realmID)
	if err != nil {
		return "", err
	}
	if realm.Type() != "github" {
		return "", fmt.Errorf("Bad realm type: %s", realm.Type())
	}

	// pull out the token (TODO: should the service know how the realm stores this?)
	session, err := database.GetServiceDB().LoadAuthSessionByUser(realm.ID(), userID)
	if err != nil {
		return "", err
	}
	ghSession, ok := session.(*realms.GithubSession)
	if !ok {
		return "", fmt.Errorf("Session is not a github session: %s", session.ID())
	}
	if ghSession.AccessToken == "" {
		return "", fmt.Errorf("Github auth session for %s has not been completed.", userID)
	}
	return ghSession.AccessToken, nil
}

func removeHook(logger *log.Entry, cli *github.Client, owner, repo, webhookEndpointURL string) {
	// Get a list of webhooks for this owner/repo and find the one which has the
	// same endpoint URL which is what github uses to determine equivalence.
	hooks, _, err := cli.Repositories.ListHooks(owner, repo, nil)
	if err != nil {
		logger.WithError(err).Print("Failed to list hooks")
		return
	}
	var hook *github.Hook
	for _, h := range hooks {
		if h.Config["url"] == nil {
			logger.Print("Ignoring nil config.url")
			continue
		}
		hookURL, ok := h.Config["url"].(string)
		if !ok {
			logger.Print("Ignoring non-string config.url")
			continue
		}
		if hookURL == webhookEndpointURL {
			hook = h
			break
		}
	}
	if hook == nil {
		return // couldn't find it
	}

	_, err = cli.Repositories.DeleteHook(owner, repo, *hook.ID)
	if err != nil {
		logger.WithError(err).Print("Failed to delete hook")
	}
}

func init() {
	types.RegisterService(func(serviceID, webhookEndpointURL string) types.Service {
		return &githubService{
			id:                 serviceID,
			webhookEndpointURL: webhookEndpointURL,
		}
	})
}
<commit_msg>Don't spam github with webhook creation requests which uses up our request rate limits<commit_after>package services

import (
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/google/go-github/github"
	"github.com/matrix-org/go-neb/database"
	"github.com/matrix-org/go-neb/matrix"
	"github.com/matrix-org/go-neb/plugin"
	"github.com/matrix-org/go-neb/realms/github"
	"github.com/matrix-org/go-neb/services/github/client"
	"github.com/matrix-org/go-neb/services/github/webhook"
	"github.com/matrix-org/go-neb/types"
	"net/http"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

// Matches alphanumeric then a /, then more alphanumeric then a #, then a number.
// E.g. owner/repo#11 (issue/PR numbers) - Captured groups for owner/repo/number
var ownerRepoIssueRegex = regexp.MustCompile("([A-z0-9-_]+)/([A-z0-9-_]+)#([0-9]+)")

type githubService struct {
	id                 string
	webhookEndpointURL string
	BotUserID          string
	ClientUserID       string
	RealmID            string
	SecretToken        string
	Rooms              map[string]struct { // room_id => {}
		Repos map[string]struct { // owner/repo => { events: ["push","issue","pull_request"] }
			Events []string
		}
	}
}

func (s *githubService) ServiceUserID() string { return s.BotUserID }
func (s *githubService) ServiceID() string     { return s.id }
func (s *githubService) ServiceType() string   { return "github" }
func (s *githubService) RoomIDs() []string {
	var keys []string
	for k := range s.Rooms {
		keys = append(keys, k)
	}
	return keys
}

func (s *githubService) cmdGithubCreate(roomID, userID string, args []string) (interface{}, error) {
	cli := s.githubClientFor(userID, false)
	if cli == nil {
		r, err := database.GetServiceDB().LoadAuthRealm(s.RealmID)
		if err != nil {
			return nil, err
		}
		ghRealm, ok := r.(*realms.GithubRealm)
		if !ok {
			return nil, fmt.Errorf("Failed to cast realm %s into a GithubRealm", s.RealmID)
		}
		return matrix.StarterLinkMessage{
			Body: "You need to OAuth with Github before you can create issues.",
			Link: ghRealm.StarterLink,
		}, nil
	}

	if len(args) < 2 {
		return &matrix.TextMessage{"m.notice",
			`Usage: !github create owner/repo "issue title" "description"`}, nil
	}

	var (
		ownerRepo string
		title     *string
		desc      *string
	)
	ownerRepo = args[0]
	o := strings.Split(ownerRepo, "/")
	if len(o) != 2 {
		return &matrix.TextMessage{"m.notice",
			`Usage: !github create owner/repo "issue title" "description"`}, nil
	}

	if len(args) == 2 {
		title = &args[1]
	} else if len(args) == 3 {
		title = &args[1]
		desc = &args[2]
	} else { // > 3 args is probably a title without quote marks
		joinedTitle := strings.Join(args[1:], " ")
		title = &joinedTitle
	}

	issue, res, err := cli.Issues.Create(o[0], o[1], &github.IssueRequest{
		Title: title,
		Body:  desc,
	})
	if err != nil {
		log.WithField("err", err).Print("Failed to create issue")
		return nil, fmt.Errorf("Failed to create issue. HTTP %d", res.StatusCode)
	}

	return matrix.TextMessage{"m.notice", fmt.Sprintf("Created issue: %s", *issue.HTMLURL)}, nil
}

func (s *githubService) expandIssue(roomID, userID string, matchingGroups []string) interface{} {
	// matchingGroups => ["foo/bar#11", "foo", "bar", "11"]
	if len(matchingGroups) != 4 {
		log.WithField("groups", matchingGroups).Print("Unexpected number of groups")
		return nil
	}
	num, err := strconv.Atoi(matchingGroups[3])
	if err != nil {
		log.WithField("issue_number", matchingGroups[3]).Print("Bad issue number")
		return nil
	}
	owner := matchingGroups[1]
	repo := matchingGroups[2]

	cli := s.githubClientFor(userID, true)

	i, _, err := cli.Issues.Get(owner, repo, num)
	if err != nil {
		log.WithError(err).WithFields(log.Fields{
			"owner":  owner,
			"repo":   repo,
			"number": num,
		}).Print("Failed to fetch issue")
		return nil
	}

	return &matrix.TextMessage{
		"m.notice",
		fmt.Sprintf("%s : %s", *i.HTMLURL, *i.Title),
	}
}

func (s *githubService) Plugin(roomID string) plugin.Plugin {
	return plugin.Plugin{
		Commands: []plugin.Command{
			plugin.Command{
				Path: []string{"github", "create"},
				Command: func(roomID, userID string, args []string) (interface{}, error) {
					return s.cmdGithubCreate(roomID, userID, args)
				},
			},
		},
		Expansions: []plugin.Expansion{
			plugin.Expansion{
				Regexp: ownerRepoIssueRegex,
				Expand: func(roomID, userID string, matchingGroups []string) interface{} {
					return s.expandIssue(roomID, userID, matchingGroups)
				},
			},
		},
	}
}
func (s *githubService) OnReceiveWebhook(w http.ResponseWriter, req *http.Request, cli *matrix.Client) {
	evType, repo, msg, err := webhook.OnReceiveRequest(req, s.SecretToken)
	if err != nil {
		w.WriteHeader(err.Code)
		return
	}

	for roomID, roomConfig := range s.Rooms {
		for ownerRepo, repoConfig := range roomConfig.Repos {
			if !strings.EqualFold(*repo.FullName, ownerRepo) {
				continue
			}

			notifyRoom := false
			for _, notifyType := range repoConfig.Events {
				if evType == notifyType {
					notifyRoom = true
					break
				}
			}
			if notifyRoom {
				log.WithFields(log.Fields{
					"type":    evType,
					"msg":     msg,
					"repo":    repo,
					"room_id": roomID,
				}).Print("Sending notification to room")
				_, e := cli.SendMessageEvent(roomID, "m.room.message", msg)
				if e != nil {
					log.WithError(e).WithField("room_id", roomID).Print(
						"Failed to send notification to room.")
				}
			}
		}
	}
	w.WriteHeader(200)
}
func (s *githubService) Register() error {
	if s.RealmID == "" || s.ClientUserID == "" || s.BotUserID == "" {
		return fmt.Errorf("RealmID, BotUserID and ClientUserID are required")
	}
	// check realm exists
	realm, err := database.GetServiceDB().LoadAuthRealm(s.RealmID)
	if err != nil {
		return err
	}
	// make sure the realm is of the type we expect
	if realm.Type() != "github" {
		return fmt.Errorf("Realm is of type '%s', not 'github'", realm.Type())
	}

	// In order to register the GH service, you must have authed with GH.
	cli := s.githubClientFor(s.ClientUserID, false)
	if cli == nil {
		return fmt.Errorf(
			"User %s does not have a Github auth session with realm %s.", s.ClientUserID, realm.ID())
	}

	return nil
}

func (s *githubService) PostRegister(oldService types.Service) {
	cli := s.githubClientFor(s.ClientUserID, false)
	if cli == nil {
		log.Errorf("PostRegister: %s does not have a github session", s.ClientUserID)
		return
	}

	if oldService != nil {
		old, ok := oldService.(*githubService)
		if !ok {
			log.Error("PostRegister: Provided old service is not of type GithubService")
			return
		}

		// Don't spam github webhook requests if we can help it.
		if sameRepos(s, old) {
			log.Print("PostRegister: old and new services have the same repo set. Nooping.")
			return
		}

		// TODO: We should be adding webhooks in Register() then removing old hooks in PostRegister()
		//
		// By doing both operations in PostRegister(), if some of the requests fail we can end up in
		// an inconsistent state. It is a lot simpler and easy to reason about this way though, so
		// for now it will do.

		// remove any existing webhooks this service created on the user's behalf
		modifyWebhooks(old, cli, true)
	}

	// make new webhooks according to service config
	modifyWebhooks(s, cli, false)
}

func modifyWebhooks(s *githubService, cli *github.Client, removeHooks bool) {
	ownerRepoSet := make(map[string]bool)
	for _, roomCfg := range s.Rooms {
		for ownerRepo := range roomCfg.Repos {
			// sanity check that it looks like 'owner/repo' as we'll split on / later
			if strings.Count(ownerRepo, "/") != 1 {
				log.WithField("owner_repo", ownerRepo).Print("Bad owner/repo value.")
				continue
			}
			ownerRepoSet[ownerRepo] = true
		}
	}

	for ownerRepo := range ownerRepoSet {
		o := strings.Split(ownerRepo, "/")
		owner := o[0]
		repo := o[1]
		logger := log.WithFields(log.Fields{
			"owner": owner,
			"repo":  repo,
		})
		if removeHooks {
			removeHook(logger, cli, owner, repo, s.webhookEndpointURL)
		} else {
			// make a hook for all GH events since we'll filter it when we receive webhook requests
			name := "web" // https://developer.github.com/v3/repos/hooks/#create-a-hook
			cfg := map[string]interface{}{
				"content_type": "json",
				"url":          s.webhookEndpointURL,
			}
			if s.SecretToken != "" {
				cfg["secret"] = s.SecretToken
			}
			events := []string{"push", "pull_request", "issues", "issue_comment", "pull_request_review_comment"}
			_, _, err := cli.Repositories.CreateHook(owner, repo, &github.Hook{
				Name:   &name,
				Config: cfg,
				Events: events,
			})
			if err != nil {
				logger.WithError(err).Print("Failed to create webhook")
				// continue as others may succeed
			}
		}
	}
}

func sameRepos(a *githubService, b *githubService) bool {
	getRepos := func(s *githubService) []string {
		r := make(map[string]bool)
		for _, roomConfig := range s.Rooms {
			for ownerRepo, _ := range roomConfig.Repos {
				r[ownerRepo] = true
			}
		}
		var rs []string
		for k, _ := range r {
			rs = append(rs, k)
		}
		return rs
	}
	aRepos := getRepos(a)
	bRepos := getRepos(b)

	if len(aRepos) != len(bRepos) {
		return false
	}

	sort.Strings(aRepos)
	sort.Strings(bRepos)
	for i := 0; i < len(aRepos); i += 1 {
		if aRepos[i] != bRepos[i] {
			return false
		}
	}
	return true
}

func (s *githubService) githubClientFor(userID string, allowUnauth bool) *github.Client {
	token, err := getTokenForUser(s.RealmID, userID)
	if err != nil {
		log.WithFields(log.Fields{
			log.ErrorKey: err,
			"user_id":    userID,
			"realm_id":   s.RealmID,
		}).Print("Failed to get token for user")
	}
	if token != "" {
		return client.New(token)
	} else if allowUnauth {
		return client.New("")
	} else {
		return nil
	}
}

func getTokenForUser(realmID, userID string) (string, error) {
	realm, err := database.GetServiceDB().LoadAuthRealm(realmID)
	if err != nil {
		return "", err
	}
	if realm.Type() != "github" {
		return "", fmt.Errorf("Bad realm type: %s", realm.Type())
	}

	// pull out the token (TODO: should the service know how the realm stores this?)
	session, err := database.GetServiceDB().LoadAuthSessionByUser(realm.ID(), userID)
	if err != nil {
		return "", err
	}
	ghSession, ok := session.(*realms.GithubSession)
	if !ok {
		return "", fmt.Errorf("Session is not a github session: %s", session.ID())
	}
	if ghSession.AccessToken == "" {
		return "", fmt.Errorf("Github auth session for %s has not been completed.", userID)
	}
	return ghSession.AccessToken, nil
}

func removeHook(logger *log.Entry, cli *github.Client, owner, repo, webhookEndpointURL string) {
	// Get a list of webhooks for this owner/repo and find the one which has the
	// same endpoint URL which is what github uses to determine equivalence.
	hooks, _, err := cli.Repositories.ListHooks(owner, repo, nil)
	if err != nil {
		logger.WithError(err).Print("Failed to list hooks")
		return
	}
	var hook *github.Hook
	for _, h := range hooks {
		if h.Config["url"] == nil {
			logger.Print("Ignoring nil config.url")
			continue
		}
		hookURL, ok := h.Config["url"].(string)
		if !ok {
			logger.Print("Ignoring non-string config.url")
			continue
		}
		if hookURL == webhookEndpointURL {
			hook = h
			break
		}
	}
	if hook == nil {
		return // couldn't find it
	}

	_, err = cli.Repositories.DeleteHook(owner, repo, *hook.ID)
	if err != nil {
		logger.WithError(err).Print("Failed to delete hook")
	}
}

func init() {
	types.RegisterService(func(serviceID, webhookEndpointURL string) types.Service {
		return &githubService{
			id:                 serviceID,
			webhookEndpointURL: webhookEndpointURL,
		}
	})
}
<|endoftext|>


<commit_before>// Copyright 2017 Authors of Cilium
//
// 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 RuntimeTest

import (
	"github.com/cilium/cilium/test/helpers"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	log "github.com/sirupsen/logrus"
)

var _ = Describe("RuntimeChaosMonkey", func() {

	var initialized bool
	var networkName string = "cilium-net"
	var netperfImage string = "tgraf/netperf"
	var logger *log.Entry
	var docker *helpers.Docker
	var cilium *helpers.Cilium

	initialize := func() {
		if initialized == true {
			return
		}
		logger = log.WithFields(log.Fields{"testName": "RuntimeChaosMonkey"})
		logger.Info("Starting")
		docker, cilium = helpers.CreateNewRuntimeHelper("runtime", logger)
		docker.NetworkCreate(networkName, "")
		initialized = true
	}

	waitForCilium := func() {
		err := cilium.WaitUntilReady(100)
		Expect(err).Should(BeNil())

		status := cilium.EndpointWaitUntilReady()
		Expect(status).Should(BeTrue())

	}

	BeforeEach(func() {
		initialize()
		docker.ContainerCreate("client", netperfImage, networkName, "-l id.client")
		docker.ContainerCreate("server", netperfImage, networkName, "-l id.server")

	})

	AfterEach(func() {
		docker.ContainerRm("client")
		docker.ContainerRm("server")
	})

	It("Endpoint recovery on restart", func() {
		originalIps := cilium.Node.Exec(`
		curl -s --unix-socket /var/run/cilium/cilium.sock \
		http://localhost/v1beta/healthz/ | jq ".ipam.ipv4|length"`)

		res := cilium.Node.Exec("sudo systemctl restart cilium")
		res.ExpectSuccess()

		waitForCilium()

		ips := cilium.Node.Exec(`
		curl -s --unix-socket /var/run/cilium/cilium.sock \
		http://localhost/v1beta/healthz/ | jq ".ipam.ipv4|length"`)
		Expect(originalIps.Output()).To(Equal(ips.Output()))

	}, 300)

	It("removing leftover Cilium interfaces", func() {
		originalLinks, err := docker.Node.Exec("sudo ip link show | wc -l").IntOutput()
		Expect(err).Should(BeNil())

		_ = docker.Node.Exec("sudo ip link add lxc12345 type veth peer name tmp54321")

		res := cilium.Node.Exec("sudo systemctl restart cilium")
		res.ExpectSuccess()

		waitForCilium()

		status := docker.Node.Exec("sudo ip link show lxc12345")
		status.ExpectFail("leftover interface were not properly cleaned up")

		links, err := docker.Node.Exec("sudo ip link show | wc -l").IntOutput()
		Expect(links).Should(Equal(originalLinks),
			"Some network interfaces were accidentally removed!")
	}, 300)
})
<commit_msg>test/runtime: make chaos test more reliable<commit_after>// Copyright 2017 Authors of Cilium
//
// 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 RuntimeTest

import (
	"github.com/cilium/cilium/test/helpers"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	log "github.com/sirupsen/logrus"
)

var _ = Describe("RuntimeChaosMonkey", func() {

	var initialized bool
	var networkName string = "cilium-net"
	var netperfImage string = "tgraf/netperf"
	var logger *log.Entry
	var docker *helpers.Docker
	var cilium *helpers.Cilium

	initialize := func() {
		if initialized == true {
			return
		}
		logger = log.WithFields(log.Fields{"testName": "RuntimeChaosMonkey"})
		logger.Info("Starting")
		docker, cilium = helpers.CreateNewRuntimeHelper("runtime", logger)
		docker.NetworkCreate(networkName, "")
		initialized = true
	}

	waitForCilium := func() {
		err := cilium.WaitUntilReady(100)
		Expect(err).Should(BeNil())

		status := cilium.EndpointWaitUntilReady()
		Expect(status).Should(BeTrue())

	}

	BeforeEach(func() {
		initialize()
		docker.ContainerCreate("client", netperfImage, networkName, "-l id.client")
		docker.ContainerCreate("server", netperfImage, networkName, "-l id.server")

		areEndpointsReady := cilium.EndpointWaitUntilReady()
		Expect(areEndpointsReady).Should(BeTrue())
	})

	AfterEach(func() {
		if CurrentGinkgoTestDescription().Failed {
			cilium.ReportFailed()
		}
		docker.ContainerRm("client")
		docker.ContainerRm("server")
	})

	It("Endpoint recovery on restart", func() {
		originalIps := cilium.Node.Exec(`
		curl -s --unix-socket /var/run/cilium/cilium.sock \
		http://localhost/v1beta/healthz/ | jq ".ipam.ipv4|length"`)

		res := cilium.Node.Exec("sudo systemctl restart cilium")
		res.ExpectSuccess()

		waitForCilium()

		ips := cilium.Node.Exec(`
		curl -s --unix-socket /var/run/cilium/cilium.sock \
		http://localhost/v1beta/healthz/ | jq ".ipam.ipv4|length"`)
		Expect(originalIps.Output().String()).To(Equal(ips.Output().String()))

	}, 300)

	It("removing leftover Cilium interfaces", func() {
		originalLinks, err := docker.Node.Exec("sudo ip link show | wc -l").IntOutput()
		Expect(err).Should(BeNil())

		_ = docker.Node.Exec("sudo ip link add lxc12345 type veth peer name tmp54321")

		res := cilium.Node.Exec("sudo systemctl restart cilium")
		res.ExpectSuccess()

		waitForCilium()

		status := docker.Node.Exec("sudo ip link show lxc12345")
		status.ExpectFail("leftover interface were not properly cleaned up")

		links, err := docker.Node.Exec("sudo ip link show | wc -l").IntOutput()
		Expect(links).Should(Equal(originalLinks),
			"Some network interfaces were accidentally removed!")
	}, 300)
})
<|endoftext|>


<commit_before>package net
<commit_msg>tests: Dial + Listen<commit_after>package net

import (
	"bytes"
	"net"
	"sync"
	"testing"

	ma "github.com/jbenet/go-multiaddr"
)

func newMultiaddr(t *testing.T, m string) ma.Multiaddr {
	maddr, err := ma.NewMultiaddr(m)
	if err != nil {
		t.Fatalf("failed to construct multiaddr: %s", m)
	}
	return maddr
}

func TestDial(t *testing.T) {

	listener, err := net.Listen("tcp", "127.0.0.1:4321")
	if err != nil {
		t.Fatal("failed to listen")
	}

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {

		cB, err := listener.Accept()
		if err != nil {
			t.Fatal("failed to accept")
		}

		// echo out
		buf := make([]byte, 1024)
		for {
			_, err := cB.Read(buf)
			if err != nil {
				break
			}
			cB.Write(buf)
		}

		wg.Done()
	}()

	maddr := newMultiaddr(t, "/ip4/127.0.0.1/tcp/4321")
	cA, err := Dial(maddr)
	if err != nil {
		t.Fatal("failed to dial")
	}

	buf := make([]byte, 1024)
	if _, err := cA.Write([]byte("beep boop")); err != nil {
		t.Fatal("failed to write:", err)
	}

	if _, err := cA.Read(buf); err != nil {
		t.Fatal("failed to read:", buf, err)
	}

	if !bytes.Equal(buf[:9], []byte("beep boop")) {
		t.Fatal("failed to echo:", buf)
	}

	maddr2 := cA.RemoteMultiaddr()
	if !maddr2.Equal(maddr) {
		t.Fatal("remote multiaddr not equal:", maddr, maddr2)
	}

	cA.Close()
	wg.Wait()
}

func TestListen(t *testing.T) {

	maddr := newMultiaddr(t, "/ip4/127.0.0.1/tcp/4322")
	listener, err := Listen(maddr)
	if err != nil {
		t.Fatal("failed to listen")
	}

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {

		cB, err := listener.Accept()
		if err != nil {
			t.Fatal("failed to accept")
		}

		if !cB.LocalMultiaddr().Equal(maddr) {
			t.Fatal("local multiaddr not equal:", maddr, cB.LocalMultiaddr())
		}

		// echo out
		buf := make([]byte, 1024)
		for {
			_, err := cB.Read(buf)
			if err != nil {
				break
			}
			cB.Write(buf)
		}

		wg.Done()
	}()

	cA, err := net.Dial("tcp", "127.0.0.1:4322")
	if err != nil {
		t.Fatal("failed to dial")
	}

	buf := make([]byte, 1024)
	if _, err := cA.Write([]byte("beep boop")); err != nil {
		t.Fatal("failed to write:", err)
	}

	if _, err := cA.Read(buf); err != nil {
		t.Fatal("failed to read:", buf, err)
	}

	if !bytes.Equal(buf[:9], []byte("beep boop")) {
		t.Fatal("failed to echo:", buf)
	}

	maddr2, err := FromNetAddr(cA.RemoteAddr())
	if err != nil {
		t.Fatal("failed to convert", err)
	}
	if !maddr2.Equal(maddr) {
		t.Fatal("remote multiaddr not equal:", maddr, maddr2)
	}

	cA.Close()
	wg.Wait()
}

func TestListenAndDial(t *testing.T) {

	maddr := newMultiaddr(t, "/ip4/127.0.0.1/tcp/4323")
	listener, err := Listen(maddr)
	if err != nil {
		t.Fatal("failed to listen")
	}

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {

		cB, err := listener.Accept()
		if err != nil {
			t.Fatal("failed to accept")
		}

		if !cB.LocalMultiaddr().Equal(maddr) {
			t.Fatal("local multiaddr not equal:", maddr, cB.LocalMultiaddr())
		}

		// echo out
		buf := make([]byte, 1024)
		for {
			_, err := cB.Read(buf)
			if err != nil {
				break
			}
			cB.Write(buf)
		}

		wg.Done()
	}()

	cA, err := Dial(newMultiaddr(t, "/ip4/127.0.0.1/tcp/4323"))
	if err != nil {
		t.Fatal("failed to dial")
	}

	buf := make([]byte, 1024)
	if _, err := cA.Write([]byte("beep boop")); err != nil {
		t.Fatal("failed to write:", err)
	}

	if _, err := cA.Read(buf); err != nil {
		t.Fatal("failed to read:", buf, err)
	}

	if !bytes.Equal(buf[:9], []byte("beep boop")) {
		t.Fatal("failed to echo:", buf)
	}

	maddr2 := cA.RemoteMultiaddr()
	if !maddr2.Equal(maddr) {
		t.Fatal("remote multiaddr not equal:", maddr, maddr2)
	}

	cA.Close()
	wg.Wait()
}
<|endoftext|>


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

package main

import (
	"context"
	"flag"
	"os"
	"path/filepath"
	"strings"

	"cloud.google.com/go/storage"
	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/skerr"
	"go.skia.org/infra/task_driver/go/lib/auth_steps"
	"go.skia.org/infra/task_driver/go/lib/docker"
	"go.skia.org/infra/task_driver/go/lib/os_steps"
	"go.skia.org/infra/task_driver/go/td"
)

const dockerImage = "gcr.io/skia-public/canvaskit-emsdk:2.0.20_v1"
const innerBuildScript = "/SRC/infra/canvaskit/build_gmtests.sh"

func main() {
	var (
		// Required properties for this task.
		outPath   = flag.String("out_path", "", "The directory to put the built wasm/js code.")
		projectID = flag.String("project_id", "", "ID of the Google Cloud project.")
		skiaPath  = flag.String("skia_path", "", "Path to skia repo root.")
		taskID    = flag.String("task_id", "", "task id this data was generated on")
		taskName  = flag.String("task_name", "", "Name of the task.")
		workPath  = flag.String("work_path", "", "The directory to use to store temporary files (e.g. docker build)")

		// Debugging flags.
		local       = flag.Bool("local", false, "True if running locally (as opposed to on the bots)")
		outputSteps = flag.String("o", "", "If provided, dump a JSON blob of step data to the given file. Prints to stdout if '-' is given.")
	)

	// Setup.
	ctx := td.StartRun(projectID, taskID, taskName, outputSteps, local)
	defer td.EndRun(ctx)

	outAbsPath := td.MustGetAbsolutePathOfFlag(ctx, *outPath, "out_path")
	skiaAbsPath := td.MustGetAbsolutePathOfFlag(ctx, *skiaPath, "skia_path")
	workAbsPath := td.MustGetAbsolutePathOfFlag(ctx, *workPath, "work_path")

	if err := os_steps.MkdirAll(ctx, workAbsPath); err != nil {
		td.Fatal(ctx, err)
	}
	if err := os_steps.MkdirAll(ctx, outAbsPath); err != nil {
		td.Fatal(ctx, err)
	}

	doc, err := setupDocker(ctx, *local)
	if err != nil {
		td.Fatal(ctx, err)
	}
	defer doc.Cleanup(ctx)

	command := []string{innerBuildScript}
	volumes := []string{skiaAbsPath + ":/SRC", workAbsPath + ":/OUT"}

	if err := doc.Run(ctx, dockerImage, command, volumes, nil); err != nil {
		td.Fatal(ctx, err)
	}

	if err := extractOutput(ctx, workAbsPath, outAbsPath); err != nil {
		td.Fatal(ctx, err)
	}
}

func setupDocker(ctx context.Context, isLocal bool) (*docker.Docker, error) {
	ctx = td.StartStep(ctx, td.Props("setup docker").Infra())
	defer td.EndStep(ctx)
	// Create token source with scope for cloud registry (storage).
	ts, err := auth_steps.Init(ctx, isLocal, auth.SCOPE_USERINFO_EMAIL, storage.ScopeReadOnly)
	if err != nil {
		return nil, td.FailStep(ctx, err)
	}

	return docker.New(ctx, ts)
}

func extractOutput(ctx context.Context, workDir, outAbsPath string) error {
	ctx = td.StartStep(ctx, td.Props("copy compiled JS and wasm into output"))
	defer td.EndStep(ctx)

	files, err := os_steps.ReadDir(ctx, workDir)
	if err != nil {
		return td.FailStep(ctx, skerr.Wrapf(err, "getting output from %s", workDir))
	}

	for _, f := range files {
		name := f.Name()
		if strings.Contains(name, "wasm_gm_tests") {
			oldFile := filepath.Join(workDir, name)
			newFile := filepath.Join(outAbsPath, name)
			if err := os.Rename(oldFile, newFile); err != nil {
				return td.FailStep(ctx, skerr.Wrapf(err, "copying %s to %s", oldFile, newFile))
			}
		}
	}
	return nil
}
<commit_msg>[canvaskit] Rollback gm tests compile version<commit_after>// Copyright 2020 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package main

import (
	"context"
	"flag"
	"os"
	"path/filepath"
	"strings"

	"cloud.google.com/go/storage"
	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/skerr"
	"go.skia.org/infra/task_driver/go/lib/auth_steps"
	"go.skia.org/infra/task_driver/go/lib/docker"
	"go.skia.org/infra/task_driver/go/lib/os_steps"
	"go.skia.org/infra/task_driver/go/td"
)

const dockerImage = "gcr.io/skia-public/canvaskit-emsdk:2.0.10_v2"
const innerBuildScript = "/SRC/infra/canvaskit/build_gmtests.sh"

func main() {
	var (
		// Required properties for this task.
		outPath   = flag.String("out_path", "", "The directory to put the built wasm/js code.")
		projectID = flag.String("project_id", "", "ID of the Google Cloud project.")
		skiaPath  = flag.String("skia_path", "", "Path to skia repo root.")
		taskID    = flag.String("task_id", "", "task id this data was generated on")
		taskName  = flag.String("task_name", "", "Name of the task.")
		workPath  = flag.String("work_path", "", "The directory to use to store temporary files (e.g. docker build)")

		// Debugging flags.
		local       = flag.Bool("local", false, "True if running locally (as opposed to on the bots)")
		outputSteps = flag.String("o", "", "If provided, dump a JSON blob of step data to the given file. Prints to stdout if '-' is given.")
	)

	// Setup.
	ctx := td.StartRun(projectID, taskID, taskName, outputSteps, local)
	defer td.EndRun(ctx)

	outAbsPath := td.MustGetAbsolutePathOfFlag(ctx, *outPath, "out_path")
	skiaAbsPath := td.MustGetAbsolutePathOfFlag(ctx, *skiaPath, "skia_path")
	workAbsPath := td.MustGetAbsolutePathOfFlag(ctx, *workPath, "work_path")

	if err := os_steps.MkdirAll(ctx, workAbsPath); err != nil {
		td.Fatal(ctx, err)
	}
	if err := os_steps.MkdirAll(ctx, outAbsPath); err != nil {
		td.Fatal(ctx, err)
	}

	doc, err := setupDocker(ctx, *local)
	if err != nil {
		td.Fatal(ctx, err)
	}
	defer doc.Cleanup(ctx)

	command := []string{innerBuildScript}
	volumes := []string{skiaAbsPath + ":/SRC", workAbsPath + ":/OUT"}

	if err := doc.Run(ctx, dockerImage, command, volumes, nil); err != nil {
		td.Fatal(ctx, err)
	}

	if err := extractOutput(ctx, workAbsPath, outAbsPath); err != nil {
		td.Fatal(ctx, err)
	}
}

func setupDocker(ctx context.Context, isLocal bool) (*docker.Docker, error) {
	ctx = td.StartStep(ctx, td.Props("setup docker").Infra())
	defer td.EndStep(ctx)
	// Create token source with scope for cloud registry (storage).
	ts, err := auth_steps.Init(ctx, isLocal, auth.SCOPE_USERINFO_EMAIL, storage.ScopeReadOnly)
	if err != nil {
		return nil, td.FailStep(ctx, err)
	}

	return docker.New(ctx, ts)
}

func extractOutput(ctx context.Context, workDir, outAbsPath string) error {
	ctx = td.StartStep(ctx, td.Props("copy compiled JS and wasm into output"))
	defer td.EndStep(ctx)

	files, err := os_steps.ReadDir(ctx, workDir)
	if err != nil {
		return td.FailStep(ctx, skerr.Wrapf(err, "getting output from %s", workDir))
	}

	for _, f := range files {
		name := f.Name()
		if strings.Contains(name, "wasm_gm_tests") {
			oldFile := filepath.Join(workDir, name)
			newFile := filepath.Join(outAbsPath, name)
			if err := os.Rename(oldFile, newFile); err != nil {
				return td.FailStep(ctx, skerr.Wrapf(err, "copying %s to %s", oldFile, newFile))
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package txindex

import (
	"context"

	"github.com/tendermint/tendermint/types"
	cmn "github.com/tendermint/tmlibs/common"
)

const (
	subscriber = "IndexerService"
)

// IndexerService connects event bus and transaction indexer together in order
// to index transactions coming from event bus.
type IndexerService struct {
	cmn.BaseService

	idr      TxIndexer
	eventBus *types.EventBus
}

// NewIndexerService returns a new service instance.
func NewIndexerService(idr TxIndexer, eventBus *types.EventBus) *IndexerService {
	is := &IndexerService{idr: idr, eventBus: eventBus}
	is.BaseService = *cmn.NewBaseService(nil, "IndexerService", is)
	return is
}

// OnStart implements cmn.Service by subscribing for all transactions
// and indexing them by tags.
func (is *IndexerService) OnStart() error {
	blockHeadersCh := make(chan interface{})
	if err := is.eventBus.Subscribe(context.Background(), subscriber, types.EventQueryNewBlockHeader, blockHeadersCh); err != nil {
		return err
	}

	txsCh := make(chan interface{})
	if err := is.eventBus.Subscribe(context.Background(), subscriber, types.EventQueryTx, txsCh); err != nil {
		return err
	}

	go func() {
		var numTxs, got int64
		var batch *Batch
		for {
			select {
			case e := <-blockHeadersCh:
				numTxs = e.(types.EventDataNewBlockHeader).Header.NumTxs
				batch = NewBatch(numTxs)
			case e := <-txsCh:
				if batch == nil {
					panic("Expected pubsub to send block header first, but got tx event")
				}
				txResult := e.(types.EventDataTx).TxResult
				batch.Add(&txResult)
				got++
				if numTxs == got {
					is.idr.AddBatch(batch)
					batch = nil
					got = 0
				}
			}
		}
	}()
	return nil
}

// OnStop implements cmn.Service by unsubscribing from all transactions.
func (is *IndexerService) OnStop() {
	if is.eventBus.IsRunning() {
		_ = is.eventBus.UnsubscribeAll(context.Background(), subscriber)
	}
}
<commit_msg>fix tests<commit_after>package txindex

import (
	"context"

	"github.com/tendermint/tendermint/types"
	cmn "github.com/tendermint/tmlibs/common"
)

const (
	subscriber = "IndexerService"
)

// IndexerService connects event bus and transaction indexer together in order
// to index transactions coming from event bus.
type IndexerService struct {
	cmn.BaseService

	idr      TxIndexer
	eventBus *types.EventBus
}

// NewIndexerService returns a new service instance.
func NewIndexerService(idr TxIndexer, eventBus *types.EventBus) *IndexerService {
	is := &IndexerService{idr: idr, eventBus: eventBus}
	is.BaseService = *cmn.NewBaseService(nil, "IndexerService", is)
	return is
}

// OnStart implements cmn.Service by subscribing for all transactions
// and indexing them by tags.
func (is *IndexerService) OnStart() error {
	blockHeadersCh := make(chan interface{})
	if err := is.eventBus.Subscribe(context.Background(), subscriber, types.EventQueryNewBlockHeader, blockHeadersCh); err != nil {
		return err
	}

	txsCh := make(chan interface{})
	if err := is.eventBus.Subscribe(context.Background(), subscriber, types.EventQueryTx, txsCh); err != nil {
		return err
	}

	go func() {
		var numTxs, got int64
		var batch *Batch
		for {
			select {
			case e, ok := <-blockHeadersCh:
				if !ok {
					return
				}
				numTxs = e.(types.EventDataNewBlockHeader).Header.NumTxs
				batch = NewBatch(numTxs)
			case e, ok := <-txsCh:
				if !ok {
					return
				}
				if batch == nil {
					panic("Expected pubsub to send block header first, but got tx event")
				}
				txResult := e.(types.EventDataTx).TxResult
				batch.Add(&txResult)
				got++
				if numTxs == got {
					is.idr.AddBatch(batch)
					batch = nil
					got = 0
				}
			}
		}
	}()
	return nil
}

// OnStop implements cmn.Service by unsubscribing from all transactions.
func (is *IndexerService) OnStop() {
	if is.eventBus.IsRunning() {
		_ = is.eventBus.UnsubscribeAll(context.Background(), subscriber)
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 tsuru-client authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"os"

	"github.com/tsuru/tsuru/cmd"
)

const (
	version = "0.13.2"
	header  = "Supported-Tsuru"
)

func buildManager(name string) *cmd.Manager {
	lookup := func(context *cmd.Context) error {
		command := plugin{}
		return command.Run(context, nil)
	}
	m := cmd.BuildBaseManager(name, version, header, lookup)
	m.RegisterDeprecated(&appRun{}, "run")
	m.Register(&appInfo{})
	m.Register(&appCreate{})
	m.Register(&appRemove{})
	m.Register(&unitAdd{})
	m.Register(&unitRemove{})
	m.Register(appList{})
	m.RegisterDeprecated(&appLog{}, "log")
	m.Register(&appGrant{})
	m.Register(&appRevoke{})
	m.RegisterDeprecated(&appRestart{}, "restart")
	m.RegisterDeprecated(&appStart{}, "start")
	m.RegisterDeprecated(&appStop{}, "stop")
	m.RegisterDeprecated(&cnameAdd{}, "add-cname")
	m.RegisterDeprecated(&cnameRemove{}, "remove-cname")
	m.Register(&envGet{})
	m.Register(&envSet{})
	m.Register(&envUnset{})
	m.Register(&keyAdd{})
	m.Register(&keyRemove{})
	m.Register(&keyList{})
	m.Register(serviceList{})
	m.Register(&serviceAdd{})
	m.Register(&serviceRemove{})
	m.Register(serviceDoc{})
	m.Register(serviceInfo{})
	m.Register(serviceInstanceStatus{})
	m.RegisterDeprecated(&serviceBind{}, "bind")
	m.RegisterDeprecated(&serviceUnbind{}, "unbind")
	m.Register(platformList{})
	m.Register(&pluginInstall{})
	m.Register(&pluginRemove{})
	m.Register(&pluginList{})
	m.RegisterDeprecated(&appSwap{}, "swap")
	m.RegisterDeprecated(&appDeploy{}, "deploy")
	m.Register(&planList{})
	return m
}

func main() {
	name := cmd.ExtractProgramName(os.Args[0])
	manager := buildManager(name)
	manager.Run(os.Args[1:])
}
<commit_msg>register app-set-team-owner command<commit_after>// Copyright 2014 tsuru-client authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"os"

	"github.com/tsuru/tsuru/cmd"
)

const (
	version = "0.13.2"
	header  = "Supported-Tsuru"
)

func buildManager(name string) *cmd.Manager {
	lookup := func(context *cmd.Context) error {
		command := plugin{}
		return command.Run(context, nil)
	}
	m := cmd.BuildBaseManager(name, version, header, lookup)
	m.RegisterDeprecated(&appRun{}, "run")
	m.Register(&appInfo{})
	m.Register(&appCreate{})
	m.Register(&appRemove{})
	m.Register(&unitAdd{})
	m.Register(&unitRemove{})
	m.Register(appList{})
	m.RegisterDeprecated(&appLog{}, "log")
	m.Register(&appGrant{})
	m.Register(&appRevoke{})
	m.RegisterDeprecated(&appRestart{}, "restart")
	m.RegisterDeprecated(&appStart{}, "start")
	m.RegisterDeprecated(&appStop{}, "stop")
	m.RegisterDeprecated(&cnameAdd{}, "add-cname")
	m.RegisterDeprecated(&cnameRemove{}, "remove-cname")
	m.Register(&envGet{})
	m.Register(&envSet{})
	m.Register(&envUnset{})
	m.Register(&keyAdd{})
	m.Register(&keyRemove{})
	m.Register(&keyList{})
	m.Register(serviceList{})
	m.Register(&serviceAdd{})
	m.Register(&serviceRemove{})
	m.Register(serviceDoc{})
	m.Register(serviceInfo{})
	m.Register(serviceInstanceStatus{})
	m.RegisterDeprecated(&serviceBind{}, "bind")
	m.RegisterDeprecated(&serviceUnbind{}, "unbind")
	m.Register(platformList{})
	m.Register(&pluginInstall{})
	m.Register(&pluginRemove{})
	m.Register(&pluginList{})
	m.RegisterDeprecated(&appSwap{}, "swap")
	m.RegisterDeprecated(&appDeploy{}, "deploy")
	m.Register(&planList{})
	m.Register(&SetTeamOwner{})
	return m
}

func main() {
	name := cmd.ExtractProgramName(os.Args[0])
	manager := buildManager(name)
	manager.Run(os.Args[1:])
}
<|endoftext|>


<commit_before>// vi: set noet ts=4 sw=4 ft=go tw=79:
/* ttrss-tool manipulates tiny-tiny-rss subscriptions. */
package main

import (
	"flag"
	"fmt"
	"net/http"
	"os"
	"strings"
)

var _ = http.StatusContinue

// Exit Codes
const (
	EX_USAGE = 64
)

// General Flags
var (
	flagAddr string
	flagUser string
	flagPass string
)

type Cmd struct {
	Flag flag.FlagSet
	Run  func(cmd *Cmd, args []string)
}

func runList(cmd *Cmd, args []string) {
	fmt.Println("RUNNING LIST")
}

var cmds = map[string]Cmd{
	"ls": {flag.FlagSet{}, runList},
}

func init() {
	flag.StringVar(&flagAddr, "addr", "",
		"address (example: https://example.com/tt-rss/)")
	flag.StringVar(&flagUser, "user", "admin", "user to connect as")
	flag.StringVar(&flagPass, "pass", "password", "password to use")

	if !strings.HasPrefix(flagAddr, "http") {
		fmt.Fprintf(os.Stderr,
			"%s: error: address %q must start with \"http\"\n",
			os.Args[0], flagAddr)
		os.Exit(EX_USAGE)
	}

	for k, v := range cmds {
		v.Flag.Init(k, flag.PanicOnError)
	}
}

func main() {
	flag.Parse()

	if flag.NArg() < 1 {
		flag.Usage()
		os.Exit(EX_USAGE)
	}

	fmt.Println(flagAddr, flagUser, flagPass, flag.Args())
}
<commit_msg>checks flag value after parsing<commit_after>// vi: set noet ts=4 sw=4 ft=go tw=79:
/* ttrss-tool manipulates tiny-tiny-rss subscriptions. */
package main

import (
	"flag"
	"fmt"
	"net/http"
	"os"
	"strings"
)

var _ = http.StatusContinue

// Exit Codes
const (
	EX_USAGE = 64
)

// General Flags
var (
	flagAddr string
	flagUser string
	flagPass string
)

type Cmd struct {
	Flag flag.FlagSet
	Run  func(cmd *Cmd, args []string)
}

func runList(cmd *Cmd, args []string) {
	fmt.Println("RUNNING LIST")
}

var cmds = map[string]Cmd{
	"ls": {flag.FlagSet{}, runList},
}

func init() {
	flag.StringVar(&flagAddr, "addr", "",
		"address (example: https://example.com/tt-rss/)")
	flag.StringVar(&flagUser, "user", "admin", "user to connect as")
	flag.StringVar(&flagPass, "pass", "password", "password to use")

	for name, cmd := range cmds {
		cmd.Flag.Init(name, flag.PanicOnError)
	}
}

func main() {
	flag.Parse()

	if flag.NArg() < 1 {
		flag.Usage()
		os.Exit(EX_USAGE)
	}

	if !strings.HasPrefix(flagAddr, "http") {
		fmt.Fprintf(os.Stderr,
			"%s: error: address %q must start with \"http\"\n",
			os.Args[0], flagAddr)
		os.Exit(EX_USAGE)
	}
	fmt.Println(flagAddr, flagUser, flagPass, flag.Args())
}
<|endoftext|>


<commit_before>package statsdclienttest

import (
	"sync"
	"time"
)

// Any type used for reporting errors. This will usually be your testing.T variable
// in tests.
type Testable interface {
	Errorf(format string, args ...interface{})
}

// StatsCommand represents a stat that was logged
type StatsCommand struct {
	Operation  string
	Stat       string
	Value      int
	SampleRate float64
}

// A stat logger used for tests
type StatsClient struct {
	// The list of stat commands that have been issued to the stat logger
	commands map[StatsCommand]int

	// The accumulated values of each stat
	Values map[string]int

	// Whether or not the logger was closed
	Closed bool

	// mutex for making updates to the underlying map atomic
	mutex sync.RWMutex
}

func (m *StatsClient) SetPrefix(prefix string) {

}

func (m *StatsClient) Unique(stat string, value int, sampleRate float64) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.commands[StatsCommand{"Unique", stat, value, sampleRate}] += 1
	m.Values[stat] = value
	return nil
}

func (m *StatsClient) Increment(stat string, delta int, sampleRate float64) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.commands[StatsCommand{"Increment", stat, delta, sampleRate}] += 1
	m.Values[stat] += delta
	return nil
}

func (m *StatsClient) Decrement(stat string, delta int, sampleRate float64) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.commands[StatsCommand{"Decrement", stat, delta, sampleRate}] += 1
	m.Values[stat] -= delta
	return nil
}

func (m *StatsClient) Gauge(stat string, value int, sampleRate float64) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.commands[StatsCommand{"Gauge", stat, value, sampleRate}] += 1
	m.Values[stat] = value
	return nil
}

func (m *StatsClient) Duration(stat string, duration time.Duration, sampleRate float64) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.commands[StatsCommand{"Duration", stat, int(duration), sampleRate}] += 1
	m.Values[stat] = int(duration / time.Millisecond)
	return nil
}

func (m *StatsClient) Close() error {
	m.Closed = true
	return nil
}

// AssertStat asserts that a given stat is in the list of logged stats, then removes it
func (m *StatsClient) AssertStat(t Testable, stat StatsCommand) {
	m.AssertLoggedN(t, stat, 1)
	delete(m.commands, stat)
}

// AssertValue asserts that the value of the stat with the given stat matches the given value.
// If the stat has not been logged, the test will fail.
func (m *StatsClient) AssertValue(t Testable, stat string, value int) {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	actualValue, ok := m.Values[stat]
	if !ok {
		t.Errorf("expected stat %q to be logged", stat)
		return
	}
	if actualValue != value {
		t.Errorf("got %d for stat %q, expected %d", actualValue, stat, value)
	}
}

// AssertLogged asserts that any stat with the given stat was logged
func (m *StatsClient) AssertLogged(t Testable, stat string) {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	if _, ok := m.Values[stat]; !ok {
		t.Errorf("expected stat %q to be logged", stat)
	}
}

func (m *StatsClient) AssertLoggedN(t Testable, stat StatsCommand, n int) {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	if m.commands[stat] != n {
		t.Errorf("stat %q logged %d times, expected %d times", stat, m.commands[stat], n)
	}
}

// NewStatsClient retur
func NewStatsClient() *StatsClient {
	return &StatsClient{
		commands: make(map[StatsCommand]int),
		Values:   make(map[string]int),
	}
}
<commit_msg>Update off latest master, fix go vet error<commit_after>package statsdclienttest

import (
	"sync"
	"time"
)

// Any type used for reporting errors. This will usually be your testing.T variable
// in tests.
type Testable interface {
	Errorf(format string, args ...interface{})
}

// StatsCommand represents a stat that was logged
type StatsCommand struct {
	Operation  string
	Stat       string
	Value      int
	SampleRate float64
}

// A stat logger used for tests
type StatsClient struct {
	// The list of stat commands that have been issued to the stat logger
	commands map[StatsCommand]int

	// The accumulated values of each stat
	Values map[string]int

	// Whether or not the logger was closed
	Closed bool

	// mutex for making updates to the underlying map atomic
	mutex sync.RWMutex
}

func (m *StatsClient) SetPrefix(prefix string) {

}

func (m *StatsClient) Unique(stat string, value int, sampleRate float64) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.commands[StatsCommand{"Unique", stat, value, sampleRate}] += 1
	m.Values[stat] = value
	return nil
}

func (m *StatsClient) Increment(stat string, delta int, sampleRate float64) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.commands[StatsCommand{"Increment", stat, delta, sampleRate}] += 1
	m.Values[stat] += delta
	return nil
}

func (m *StatsClient) Decrement(stat string, delta int, sampleRate float64) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.commands[StatsCommand{"Decrement", stat, delta, sampleRate}] += 1
	m.Values[stat] -= delta
	return nil
}

func (m *StatsClient) Gauge(stat string, value int, sampleRate float64) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.commands[StatsCommand{"Gauge", stat, value, sampleRate}] += 1
	m.Values[stat] = value
	return nil
}

func (m *StatsClient) Duration(stat string, duration time.Duration, sampleRate float64) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.commands[StatsCommand{"Duration", stat, int(duration), sampleRate}] += 1
	m.Values[stat] = int(duration / time.Millisecond)
	return nil
}

func (m *StatsClient) Close() error {
	m.Closed = true
	return nil
}

// AssertStat asserts that a given stat is in the list of logged stats, then removes it
func (m *StatsClient) AssertStat(t Testable, stat StatsCommand) {
	m.AssertLoggedN(t, stat, 1)
	delete(m.commands, stat)
}

// AssertValue asserts that the value of the stat with the given stat matches the given value.
// If the stat has not been logged, the test will fail.
func (m *StatsClient) AssertValue(t Testable, stat string, value int) {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	actualValue, ok := m.Values[stat]
	if !ok {
		t.Errorf("expected stat %q to be logged", stat)
		return
	}
	if actualValue != value {
		t.Errorf("got %d for stat %q, expected %d", actualValue, stat, value)
	}
}

// AssertLogged asserts that any stat with the given stat was logged
func (m *StatsClient) AssertLogged(t Testable, stat string) {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	if _, ok := m.Values[stat]; !ok {
		t.Errorf("expected stat %q to be logged", stat)
	}
}

func (m *StatsClient) AssertLoggedN(t Testable, stat StatsCommand, n int) {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	if m.commands[stat] != n {
		t.Errorf("stat %+v logged %d times, expected %d times", stat, m.commands[stat], n)
	}
}

// NewStatsClient retur
func NewStatsClient() *StatsClient {
	return &StatsClient{
		commands: make(map[StatsCommand]int),
		Values:   make(map[string]int),
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 fatedier, fatedier@gmail.com
//
// 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 server

import (
	"fmt"
	"time"

	"github.com/fatedier/frp/assets"
	"github.com/fatedier/frp/models/config"
	"github.com/fatedier/frp/models/msg"
	"github.com/fatedier/frp/utils/log"
	frpNet "github.com/fatedier/frp/utils/net"
	"github.com/fatedier/frp/utils/util"
	"github.com/fatedier/frp/utils/version"
	"github.com/fatedier/frp/utils/vhost"

	"github.com/xtaci/smux"
)

const (
	connReadTimeout time.Duration = 10 * time.Second
)

var ServerService *Service

// Server service.
type Service struct {
	// Accept connections from client.
	listener frpNet.Listener

	// Accept connections using kcp.
	kcpListener frpNet.Listener

	// For http proxies, route requests to different clients by hostname and other infomation.
	VhostHttpMuxer *vhost.HttpMuxer

	// For https proxies, route requests to different clients by hostname and other infomation.
	VhostHttpsMuxer *vhost.HttpsMuxer

	// Manage all controllers.
	ctlManager *ControlManager

	// Manage all proxies.
	pxyManager *ProxyManager

	// Manage all visitor listeners.
	visitorManager *VisitorManager

	// Controller for nat hole connections.
	natHoleController *NatHoleController
}

func NewService() (svr *Service, err error) {
	svr = &Service{
		ctlManager:     NewControlManager(),
		pxyManager:     NewProxyManager(),
		visitorManager: NewVisitorManager(),
	}
	cfg := config.ServerCommonCfg

	// Init assets.
	err = assets.Load(cfg.AssetsDir)
	if err != nil {
		err = fmt.Errorf("Load assets error: %v", err)
		return
	}

	// Listen for accepting connections from client.
	svr.listener, err = frpNet.ListenTcp(cfg.BindAddr, cfg.BindPort)
	if err != nil {
		err = fmt.Errorf("Create server listener error, %v", err)
		return
	}
	log.Info("frps tcp listen on %s:%d", cfg.BindAddr, cfg.BindPort)

	// Listen for accepting connections from client using kcp protocol.
	if cfg.KcpBindPort > 0 {
		svr.kcpListener, err = frpNet.ListenKcp(cfg.BindAddr, cfg.KcpBindPort)
		if err != nil {
			err = fmt.Errorf("Listen on kcp address udp [%s:%d] error: %v", cfg.BindAddr, cfg.KcpBindPort, err)
			return
		}
		log.Info("frps kcp listen on udp %s:%d", cfg.BindAddr, cfg.BindPort)
	}

	// Create http vhost muxer.
	if cfg.VhostHttpPort > 0 {
		var l frpNet.Listener
		l, err = frpNet.ListenTcp(cfg.ProxyBindAddr, cfg.VhostHttpPort)
		if err != nil {
			err = fmt.Errorf("Create vhost http listener error, %v", err)
			return
		}
		svr.VhostHttpMuxer, err = vhost.NewHttpMuxer(l, 30*time.Second)
		if err != nil {
			err = fmt.Errorf("Create vhost httpMuxer error, %v", err)
			return
		}
		log.Info("http service listen on %s:%d", cfg.ProxyBindAddr, cfg.VhostHttpPort)
	}

	// Create https vhost muxer.
	if cfg.VhostHttpsPort > 0 {
		var l frpNet.Listener
		l, err = frpNet.ListenTcp(cfg.ProxyBindAddr, cfg.VhostHttpsPort)
		if err != nil {
			err = fmt.Errorf("Create vhost https listener error, %v", err)
			return
		}
		svr.VhostHttpsMuxer, err = vhost.NewHttpsMuxer(l, 30*time.Second)
		if err != nil {
			err = fmt.Errorf("Create vhost httpsMuxer error, %v", err)
			return
		}
		log.Info("https service listen on %s:%d", cfg.ProxyBindAddr, cfg.VhostHttpsPort)
	}

	// Create nat hole controller.
	if cfg.BindUdpPort > 0 {
		var nc *NatHoleController
		addr := fmt.Sprintf("%s:%d", cfg.BindAddr, cfg.BindUdpPort)
		nc, err = NewNatHoleController(addr)
		if err != nil {
			err = fmt.Errorf("Create nat hole controller error, %v", err)
			return
		}
		svr.natHoleController = nc
		log.Info("nat hole udp service listen on %s:%d", cfg.BindAddr, cfg.BindUdpPort)
	}

	// Create dashboard web server.
	if cfg.DashboardPort > 0 {
		err = RunDashboardServer(cfg.DashboardAddr, cfg.DashboardPort)
		if err != nil {
			err = fmt.Errorf("Create dashboard web server error, %v", err)
			return
		}
		log.Info("Dashboard listen on %s:%d", cfg.DashboardAddr, cfg.DashboardPort)
	}
	return
}

func (svr *Service) Run() {
	if svr.natHoleController != nil {
		go svr.natHoleController.Run()
	}
	if config.ServerCommonCfg.KcpBindPort > 0 {
		go svr.HandleListener(svr.kcpListener)
	}
	svr.HandleListener(svr.listener)

}

func (svr *Service) HandleListener(l frpNet.Listener) {
	// Listen for incoming connections from client.
	for {
		c, err := l.Accept()
		if err != nil {
			log.Warn("Listener for incoming connections from client closed")
			return
		}

		// Start a new goroutine for dealing connections.
		go func(frpConn frpNet.Conn) {
			dealFn := func(conn frpNet.Conn) {
				var rawMsg msg.Message
				conn.SetReadDeadline(time.Now().Add(connReadTimeout))
				if rawMsg, err = msg.ReadMsg(conn); err != nil {
					log.Trace("Failed to read message: %v", err)
					conn.Close()
					return
				}
				conn.SetReadDeadline(time.Time{})

				switch m := rawMsg.(type) {
				case *msg.Login:
					err = svr.RegisterControl(conn, m)
					// If login failed, send error message there.
					// Otherwise send success message in control's work goroutine.
					if err != nil {
						conn.Warn("%v", err)
						msg.WriteMsg(conn, &msg.LoginResp{
							Version: version.Full(),
							Error:   err.Error(),
						})
						conn.Close()
					}
				case *msg.NewWorkConn:
					svr.RegisterWorkConn(conn, m)
				case *msg.NewVisitorConn:
					if err = svr.RegisterVisitorConn(conn, m); err != nil {
						conn.Warn("%v", err)
						msg.WriteMsg(conn, &msg.NewVisitorConnResp{
							ProxyName: m.ProxyName,
							Error:     err.Error(),
						})
						conn.Close()
					} else {
						msg.WriteMsg(conn, &msg.NewVisitorConnResp{
							ProxyName: m.ProxyName,
							Error:     "",
						})
					}
				default:
					log.Warn("Error message type for the new connection [%s]", conn.RemoteAddr().String())
					conn.Close()
				}
			}

			if config.ServerCommonCfg.TcpMux {
				session, err := smux.Server(frpConn, nil)
				if err != nil {
					log.Warn("Failed to create mux connection: %v", err)
					frpConn.Close()
					return
				}

				for {
					stream, err := session.AcceptStream()
					if err != nil {
						log.Warn("Accept new mux stream error: %v", err)
						session.Close()
						return
					}
					wrapConn := frpNet.WrapConn(stream)
					go dealFn(wrapConn)
				}
			} else {
				dealFn(frpConn)
			}
		}(c)
	}
}

func (svr *Service) RegisterControl(ctlConn frpNet.Conn, loginMsg *msg.Login) (err error) {
	ctlConn.Info("client login info: ip [%s] version [%s] hostname [%s] os [%s] arch [%s]",
		ctlConn.RemoteAddr().String(), loginMsg.Version, loginMsg.Hostname, loginMsg.Os, loginMsg.Arch)

	// Check client version.
	if ok, msg := version.Compat(loginMsg.Version); !ok {
		err = fmt.Errorf("%s", msg)
		return
	}

	// Check auth.
	nowTime := time.Now().Unix()
	if config.ServerCommonCfg.AuthTimeout != 0 && nowTime-loginMsg.Timestamp > config.ServerCommonCfg.AuthTimeout {
		err = fmt.Errorf("authorization timeout")
		return
	}
	if util.GetAuthKey(config.ServerCommonCfg.PrivilegeToken, loginMsg.Timestamp) != loginMsg.PrivilegeKey {
		err = fmt.Errorf("authorization failed")
		return
	}

	// If client's RunId is empty, it's a new client, we just create a new controller.
	// Otherwise, we check if there is one controller has the same run id. If so, we release previous controller and start new one.
	if loginMsg.RunId == "" {
		loginMsg.RunId, err = util.RandId()
		if err != nil {
			return
		}
	}

	ctl := NewControl(svr, ctlConn, loginMsg)

	if oldCtl := svr.ctlManager.Add(loginMsg.RunId, ctl); oldCtl != nil {
		oldCtl.allShutdown.WaitDown()
	}

	ctlConn.AddLogPrefix(loginMsg.RunId)
	ctl.Start()

	// for statistics
	StatsNewClient()
	return
}

// RegisterWorkConn register a new work connection to control and proxies need it.
func (svr *Service) RegisterWorkConn(workConn frpNet.Conn, newMsg *msg.NewWorkConn) {
	ctl, exist := svr.ctlManager.GetById(newMsg.RunId)
	if !exist {
		workConn.Warn("No client control found for run id [%s]", newMsg.RunId)
		return
	}
	ctl.RegisterWorkConn(workConn)
	return
}

func (svr *Service) RegisterVisitorConn(visitorConn frpNet.Conn, newMsg *msg.NewVisitorConn) error {
	return svr.visitorManager.NewConn(newMsg.ProxyName, visitorConn, newMsg.Timestamp, newMsg.SignKey,
		newMsg.UseEncryption, newMsg.UseCompression)
}

func (svr *Service) RegisterProxy(name string, pxy Proxy) error {
	return svr.pxyManager.Add(name, pxy)
}

func (svr *Service) DelProxy(name string) {
	svr.pxyManager.Del(name)
}
<commit_msg>fix kcp port print error<commit_after>// Copyright 2017 fatedier, fatedier@gmail.com
//
// 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 server

import (
	"fmt"
	"time"

	"github.com/fatedier/frp/assets"
	"github.com/fatedier/frp/models/config"
	"github.com/fatedier/frp/models/msg"
	"github.com/fatedier/frp/utils/log"
	frpNet "github.com/fatedier/frp/utils/net"
	"github.com/fatedier/frp/utils/util"
	"github.com/fatedier/frp/utils/version"
	"github.com/fatedier/frp/utils/vhost"

	"github.com/xtaci/smux"
)

const (
	connReadTimeout time.Duration = 10 * time.Second
)

var ServerService *Service

// Server service.
type Service struct {
	// Accept connections from client.
	listener frpNet.Listener

	// Accept connections using kcp.
	kcpListener frpNet.Listener

	// For http proxies, route requests to different clients by hostname and other infomation.
	VhostHttpMuxer *vhost.HttpMuxer

	// For https proxies, route requests to different clients by hostname and other infomation.
	VhostHttpsMuxer *vhost.HttpsMuxer

	// Manage all controllers.
	ctlManager *ControlManager

	// Manage all proxies.
	pxyManager *ProxyManager

	// Manage all visitor listeners.
	visitorManager *VisitorManager

	// Controller for nat hole connections.
	natHoleController *NatHoleController
}

func NewService() (svr *Service, err error) {
	svr = &Service{
		ctlManager:     NewControlManager(),
		pxyManager:     NewProxyManager(),
		visitorManager: NewVisitorManager(),
	}
	cfg := config.ServerCommonCfg

	// Init assets.
	err = assets.Load(cfg.AssetsDir)
	if err != nil {
		err = fmt.Errorf("Load assets error: %v", err)
		return
	}

	// Listen for accepting connections from client.
	svr.listener, err = frpNet.ListenTcp(cfg.BindAddr, cfg.BindPort)
	if err != nil {
		err = fmt.Errorf("Create server listener error, %v", err)
		return
	}
	log.Info("frps tcp listen on %s:%d", cfg.BindAddr, cfg.BindPort)

	// Listen for accepting connections from client using kcp protocol.
	if cfg.KcpBindPort > 0 {
		svr.kcpListener, err = frpNet.ListenKcp(cfg.BindAddr, cfg.KcpBindPort)
		if err != nil {
			err = fmt.Errorf("Listen on kcp address udp [%s:%d] error: %v", cfg.BindAddr, cfg.KcpBindPort, err)
			return
		}
		log.Info("frps kcp listen on udp %s:%d", cfg.BindAddr, cfg.KcpBindPort)
	}

	// Create http vhost muxer.
	if cfg.VhostHttpPort > 0 {
		var l frpNet.Listener
		l, err = frpNet.ListenTcp(cfg.ProxyBindAddr, cfg.VhostHttpPort)
		if err != nil {
			err = fmt.Errorf("Create vhost http listener error, %v", err)
			return
		}
		svr.VhostHttpMuxer, err = vhost.NewHttpMuxer(l, 30*time.Second)
		if err != nil {
			err = fmt.Errorf("Create vhost httpMuxer error, %v", err)
			return
		}
		log.Info("http service listen on %s:%d", cfg.ProxyBindAddr, cfg.VhostHttpPort)
	}

	// Create https vhost muxer.
	if cfg.VhostHttpsPort > 0 {
		var l frpNet.Listener
		l, err = frpNet.ListenTcp(cfg.ProxyBindAddr, cfg.VhostHttpsPort)
		if err != nil {
			err = fmt.Errorf("Create vhost https listener error, %v", err)
			return
		}
		svr.VhostHttpsMuxer, err = vhost.NewHttpsMuxer(l, 30*time.Second)
		if err != nil {
			err = fmt.Errorf("Create vhost httpsMuxer error, %v", err)
			return
		}
		log.Info("https service listen on %s:%d", cfg.ProxyBindAddr, cfg.VhostHttpsPort)
	}

	// Create nat hole controller.
	if cfg.BindUdpPort > 0 {
		var nc *NatHoleController
		addr := fmt.Sprintf("%s:%d", cfg.BindAddr, cfg.BindUdpPort)
		nc, err = NewNatHoleController(addr)
		if err != nil {
			err = fmt.Errorf("Create nat hole controller error, %v", err)
			return
		}
		svr.natHoleController = nc
		log.Info("nat hole udp service listen on %s:%d", cfg.BindAddr, cfg.BindUdpPort)
	}

	// Create dashboard web server.
	if cfg.DashboardPort > 0 {
		err = RunDashboardServer(cfg.DashboardAddr, cfg.DashboardPort)
		if err != nil {
			err = fmt.Errorf("Create dashboard web server error, %v", err)
			return
		}
		log.Info("Dashboard listen on %s:%d", cfg.DashboardAddr, cfg.DashboardPort)
	}
	return
}

func (svr *Service) Run() {
	if svr.natHoleController != nil {
		go svr.natHoleController.Run()
	}
	if config.ServerCommonCfg.KcpBindPort > 0 {
		go svr.HandleListener(svr.kcpListener)
	}
	svr.HandleListener(svr.listener)

}

func (svr *Service) HandleListener(l frpNet.Listener) {
	// Listen for incoming connections from client.
	for {
		c, err := l.Accept()
		if err != nil {
			log.Warn("Listener for incoming connections from client closed")
			return
		}

		// Start a new goroutine for dealing connections.
		go func(frpConn frpNet.Conn) {
			dealFn := func(conn frpNet.Conn) {
				var rawMsg msg.Message
				conn.SetReadDeadline(time.Now().Add(connReadTimeout))
				if rawMsg, err = msg.ReadMsg(conn); err != nil {
					log.Trace("Failed to read message: %v", err)
					conn.Close()
					return
				}
				conn.SetReadDeadline(time.Time{})

				switch m := rawMsg.(type) {
				case *msg.Login:
					err = svr.RegisterControl(conn, m)
					// If login failed, send error message there.
					// Otherwise send success message in control's work goroutine.
					if err != nil {
						conn.Warn("%v", err)
						msg.WriteMsg(conn, &msg.LoginResp{
							Version: version.Full(),
							Error:   err.Error(),
						})
						conn.Close()
					}
				case *msg.NewWorkConn:
					svr.RegisterWorkConn(conn, m)
				case *msg.NewVisitorConn:
					if err = svr.RegisterVisitorConn(conn, m); err != nil {
						conn.Warn("%v", err)
						msg.WriteMsg(conn, &msg.NewVisitorConnResp{
							ProxyName: m.ProxyName,
							Error:     err.Error(),
						})
						conn.Close()
					} else {
						msg.WriteMsg(conn, &msg.NewVisitorConnResp{
							ProxyName: m.ProxyName,
							Error:     "",
						})
					}
				default:
					log.Warn("Error message type for the new connection [%s]", conn.RemoteAddr().String())
					conn.Close()
				}
			}

			if config.ServerCommonCfg.TcpMux {
				session, err := smux.Server(frpConn, nil)
				if err != nil {
					log.Warn("Failed to create mux connection: %v", err)
					frpConn.Close()
					return
				}

				for {
					stream, err := session.AcceptStream()
					if err != nil {
						log.Warn("Accept new mux stream error: %v", err)
						session.Close()
						return
					}
					wrapConn := frpNet.WrapConn(stream)
					go dealFn(wrapConn)
				}
			} else {
				dealFn(frpConn)
			}
		}(c)
	}
}

func (svr *Service) RegisterControl(ctlConn frpNet.Conn, loginMsg *msg.Login) (err error) {
	ctlConn.Info("client login info: ip [%s] version [%s] hostname [%s] os [%s] arch [%s]",
		ctlConn.RemoteAddr().String(), loginMsg.Version, loginMsg.Hostname, loginMsg.Os, loginMsg.Arch)

	// Check client version.
	if ok, msg := version.Compat(loginMsg.Version); !ok {
		err = fmt.Errorf("%s", msg)
		return
	}

	// Check auth.
	nowTime := time.Now().Unix()
	if config.ServerCommonCfg.AuthTimeout != 0 && nowTime-loginMsg.Timestamp > config.ServerCommonCfg.AuthTimeout {
		err = fmt.Errorf("authorization timeout")
		return
	}
	if util.GetAuthKey(config.ServerCommonCfg.PrivilegeToken, loginMsg.Timestamp) != loginMsg.PrivilegeKey {
		err = fmt.Errorf("authorization failed")
		return
	}

	// If client's RunId is empty, it's a new client, we just create a new controller.
	// Otherwise, we check if there is one controller has the same run id. If so, we release previous controller and start new one.
	if loginMsg.RunId == "" {
		loginMsg.RunId, err = util.RandId()
		if err != nil {
			return
		}
	}

	ctl := NewControl(svr, ctlConn, loginMsg)

	if oldCtl := svr.ctlManager.Add(loginMsg.RunId, ctl); oldCtl != nil {
		oldCtl.allShutdown.WaitDown()
	}

	ctlConn.AddLogPrefix(loginMsg.RunId)
	ctl.Start()

	// for statistics
	StatsNewClient()
	return
}

// RegisterWorkConn register a new work connection to control and proxies need it.
func (svr *Service) RegisterWorkConn(workConn frpNet.Conn, newMsg *msg.NewWorkConn) {
	ctl, exist := svr.ctlManager.GetById(newMsg.RunId)
	if !exist {
		workConn.Warn("No client control found for run id [%s]", newMsg.RunId)
		return
	}
	ctl.RegisterWorkConn(workConn)
	return
}

func (svr *Service) RegisterVisitorConn(visitorConn frpNet.Conn, newMsg *msg.NewVisitorConn) error {
	return svr.visitorManager.NewConn(newMsg.ProxyName, visitorConn, newMsg.Timestamp, newMsg.SignKey,
		newMsg.UseEncryption, newMsg.UseCompression)
}

func (svr *Service) RegisterProxy(name string, pxy Proxy) error {
	return svr.pxyManager.Add(name, pxy)
}

func (svr *Service) DelProxy(name string) {
	svr.pxyManager.Del(name)
}
<|endoftext|>