<commit_before>package aws

import (
	"fmt"
	"strings"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSGroupMembership_basic(t *testing.T) {
	var group iam.GetGroupOutput

	rString := acctest.RandStringFromCharSet(10, acctest.CharSetAlpha)
	configBase := fmt.Sprintf(testAccAWSGroupMemberConfig, rString, rString, rString)
	configUpdate := fmt.Sprintf(testAccAWSGroupMemberConfigUpdate, rString, rString, rString, rString)
	configUpdateDown := fmt.Sprintf(testAccAWSGroupMemberConfigUpdateDown, rString, rString, rString)

	testUser := fmt.Sprintf("test-user-%s", rString)
	testUserTwo := fmt.Sprintf("test-user-two-%s", rString)
	testUserThree := fmt.Sprintf("test-user-three-%s", rString)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSGroupMembershipDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: configBase,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSGroupMembershipExists("aws_iam_group_membership.team", &group),
					testAccCheckAWSGroupMembershipAttributes(&group, []string{testUser}),
				),
			},

			resource.TestStep{
				Config: configUpdate,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSGroupMembershipExists("aws_iam_group_membership.team", &group),
					testAccCheckAWSGroupMembershipAttributes(&group, []string{testUserTwo, testUserThree}),
				),
			},

			resource.TestStep{
				Config: configUpdateDown,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSGroupMembershipExists("aws_iam_group_membership.team", &group),
					testAccCheckAWSGroupMembershipAttributes(&group, []string{testUserThree}),
				),
			},
		},
	})
}

func testAccCheckAWSGroupMembershipDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).iamconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_iam_group_membership" {
			continue
		}

		group := rs.Primary.Attributes["group"]

		_, err := conn.GetGroup(&iam.GetGroupInput{
			GroupName: aws.String(group),
		})
		if err != nil {
			// Verify the error is what we want
			if ae, ok := err.(awserr.Error); ok && ae.Code() == "NoSuchEntity" {
				continue
			}
			return err
		}

		return fmt.Errorf("still exists")
	}

	return nil
}

func testAccCheckAWSGroupMembershipExists(n string, g *iam.GetGroupOutput) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No User name is set")
		}

		conn := testAccProvider.Meta().(*AWSClient).iamconn
		gn := rs.Primary.Attributes["group"]

		resp, err := conn.GetGroup(&iam.GetGroupInput{
			GroupName: aws.String(gn),
		})

		if err != nil {
			return fmt.Errorf("Error: Group (%s) not found", gn)
		}

		*g = *resp

		return nil
	}
}

func testAccCheckAWSGroupMembershipAttributes(group *iam.GetGroupOutput, users []string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		if !strings.Contains(*group.Group.GroupName, "test-group") {
			return fmt.Errorf("Bad group membership: expected %d, got %d", "test-group", *group.Group.GroupName)
		}

		uc := len(users)
		for _, u := range users {
			for _, gu := range group.Users {
				if u == *gu.UserName {
					uc--
				}
			}
		}

		if uc > 0 {
			return fmt.Errorf("Bad group membership count, expected (%d), but only (%d) found", len(users), uc)
		}
		return nil
	}
}

const testAccAWSGroupMemberConfig = `
resource "aws_iam_group" "group" {
	name = "test-group-%s"
	path = "/"
}

resource "aws_iam_user" "user" {
	name = "test-user-%s"
	path = "/"
}

resource "aws_iam_group_membership" "team" {
	name = "tf-testing-group-membership-%s"
	users = ["${aws_iam_user.user.name}"]
	group = "${aws_iam_group.group.name}"
}
`

const testAccAWSGroupMemberConfigUpdate = `
resource "aws_iam_group" "group" {
	name = "test-group-%s"
	path = "/"
}

resource "aws_iam_user" "user" {
	name = "test-user-%s"
	path = "/"
}

resource "aws_iam_user" "user_two" {
	name = "test-user-two-%s"
	path = "/"
}

resource "aws_iam_user" "user_three" {
	name = "test-user-three-%s"
	path = "/"
}

resource "aws_iam_group_membership" "team" {
	name = "tf-testing-group-membership-%s"
	users = [
		"${aws_iam_user.user_two.name}",
		"${aws_iam_user.user_three.name}",
	]
	group = "${aws_iam_group.group.name}"
}
`

const testAccAWSGroupMemberConfigUpdateDown = `
resource "aws_iam_group" "group" {
	name = "test-group-%s"
	path = "/"
}

resource "aws_iam_user" "user_three" {
	name = "test-user-three-%s"
	path = "/"
}

resource "aws_iam_group_membership" "team" {
	name = "tf-testing-group-membership-%s"
	users = [
		"${aws_iam_user.user_three.name}",
	]
	group = "${aws_iam_group.group.name}"
}
`
<commit_msg>Fix additional vet warnings<commit_after>package aws

import (
	"fmt"
	"strings"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSGroupMembership_basic(t *testing.T) {
	var group iam.GetGroupOutput

	rString := acctest.RandStringFromCharSet(10, acctest.CharSetAlpha)
	configBase := fmt.Sprintf(testAccAWSGroupMemberConfig, rString, rString, rString)
	configUpdate := fmt.Sprintf(testAccAWSGroupMemberConfigUpdate, rString, rString, rString, rString, rString)
	configUpdateDown := fmt.Sprintf(testAccAWSGroupMemberConfigUpdateDown, rString, rString, rString)

	testUser := fmt.Sprintf("test-user-%s", rString)
	testUserTwo := fmt.Sprintf("test-user-two-%s", rString)
	testUserThree := fmt.Sprintf("test-user-three-%s", rString)

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSGroupMembershipDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: configBase,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSGroupMembershipExists("aws_iam_group_membership.team", &group),
					testAccCheckAWSGroupMembershipAttributes(&group, []string{testUser}),
				),
			},

			resource.TestStep{
				Config: configUpdate,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSGroupMembershipExists("aws_iam_group_membership.team", &group),
					testAccCheckAWSGroupMembershipAttributes(&group, []string{testUserTwo, testUserThree}),
				),
			},

			resource.TestStep{
				Config: configUpdateDown,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSGroupMembershipExists("aws_iam_group_membership.team", &group),
					testAccCheckAWSGroupMembershipAttributes(&group, []string{testUserThree}),
				),
			},
		},
	})
}

func testAccCheckAWSGroupMembershipDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).iamconn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_iam_group_membership" {
			continue
		}

		group := rs.Primary.Attributes["group"]

		_, err := conn.GetGroup(&iam.GetGroupInput{
			GroupName: aws.String(group),
		})
		if err != nil {
			// Verify the error is what we want
			if ae, ok := err.(awserr.Error); ok && ae.Code() == "NoSuchEntity" {
				continue
			}
			return err
		}

		return fmt.Errorf("still exists")
	}

	return nil
}

func testAccCheckAWSGroupMembershipExists(n string, g *iam.GetGroupOutput) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No User name is set")
		}

		conn := testAccProvider.Meta().(*AWSClient).iamconn
		gn := rs.Primary.Attributes["group"]

		resp, err := conn.GetGroup(&iam.GetGroupInput{
			GroupName: aws.String(gn),
		})

		if err != nil {
			return fmt.Errorf("Error: Group (%s) not found", gn)
		}

		*g = *resp

		return nil
	}
}

func testAccCheckAWSGroupMembershipAttributes(group *iam.GetGroupOutput, users []string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		if !strings.Contains(*group.Group.GroupName, "test-group") {
			return fmt.Errorf("Bad group membership: expected %s, got %s", "test-group", *group.Group.GroupName)
		}

		uc := len(users)
		for _, u := range users {
			for _, gu := range group.Users {
				if u == *gu.UserName {
					uc--
				}
			}
		}

		if uc > 0 {
			return fmt.Errorf("Bad group membership count, expected (%d), but only (%d) found", len(users), uc)
		}
		return nil
	}
}

const testAccAWSGroupMemberConfig = `
resource "aws_iam_group" "group" {
	name = "test-group-%s"
	path = "/"
}

resource "aws_iam_user" "user" {
	name = "test-user-%s"
	path = "/"
}

resource "aws_iam_group_membership" "team" {
	name = "tf-testing-group-membership-%s"
	users = ["${aws_iam_user.user.name}"]
	group = "${aws_iam_group.group.name}"
}
`

const testAccAWSGroupMemberConfigUpdate = `
resource "aws_iam_group" "group" {
	name = "test-group-%s"
	path = "/"
}

resource "aws_iam_user" "user" {
	name = "test-user-%s"
	path = "/"
}

resource "aws_iam_user" "user_two" {
	name = "test-user-two-%s"
	path = "/"
}

resource "aws_iam_user" "user_three" {
	name = "test-user-three-%s"
	path = "/"
}

resource "aws_iam_group_membership" "team" {
	name = "tf-testing-group-membership-%s"
	users = [
		"${aws_iam_user.user_two.name}",
		"${aws_iam_user.user_three.name}",
	]
	group = "${aws_iam_group.group.name}"
}
`

const testAccAWSGroupMemberConfigUpdateDown = `
resource "aws_iam_group" "group" {
	name = "test-group-%s"
	path = "/"
}

resource "aws_iam_user" "user_three" {
	name = "test-user-three-%s"
	path = "/"
}

resource "aws_iam_group_membership" "team" {
	name = "tf-testing-group-membership-%s"
	users = [
		"${aws_iam_user.user_three.name}",
	]
	group = "${aws_iam_group.group.name}"
}
`
<|endoftext|>


<commit_before>/*
Copyright 2016 The Transicator Authors

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

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

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

import (
	"bytes"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"

	"time"

	log "github.com/Sirupsen/logrus"
	sqlite "github.com/mattn/go-sqlite3"
)

var sqliteTimestampFormat = sqlite.SQLiteTimestampFormats[0]

/*
WriteSqliteSnapshot is responsible for generating the SQLite format of
the snapshot, and also for streaming the file back to the
caller.
*/
func WriteSqliteSnapshot(scopes []string, db *sql.DB, w http.ResponseWriter, r *http.Request) error {

	// Create temporary directory for the temporary DB
	dirName, err := ioutil.TempDir(tempSnapshotDir, tempSnapshotPrefix)
	if err != nil {
		sendAPIError(serverError, err.Error(), w, r)
		return err
	}
	defer func() {
		os.RemoveAll(dirName)
	}()

	// Open and verify the DB
	dbFileName := path.Join(dirName, tempSnapshotName)
	tdb, err := createDatabase(dbFileName)
	if err != nil {
		sendAPIError(http.StatusInternalServerError, err.Error(), w, r)
		return err
	}
	defer tdb.Close()

	// Because of previous config, this puts us in "repeatable read" mode
	pgTx, err := db.Begin()
	if err != nil {
		sendAPIError(serverError, err.Error(), w, r)
		return err
	}
	defer pgTx.Commit()

	tables, err := enumeratePgTables(pgTx)
	if err != nil {
		sendAPIError(http.StatusInternalServerError, err.Error(), w, r)
		return err
	}

	err = writeMetadata(pgTx, tdb, tables)
	if err != nil {
		sendAPIError(http.StatusInternalServerError, err.Error(), w, r)
		return err
	}
	// put tx id into header
	row := pgTx.QueryRow("select txid_current_snapshot()")
	var txId string
	err = row.Scan(&txId)
	if err == nil {
		w.Header().Set("Transicator-Snapshot-TXID", txId)
	}

	// For each table, update the DB
	for tid, pgTable := range tables {
		if pgTable.hasSelector {
			err = makeSqliteTable(tdb, pgTable)
			if err == nil {
				err = copyData(pgTx, tdb, scopes, pgTable)
			}
			if err != nil {
				sendAPIError(serverError, err.Error(), w, r)
				return err
			}
		} else {
			log.Debugf("Skipping table %s which has no selector", tid)
		}
	}

	pgTx.Commit()

	// Close the database. Checkpoint the WAL so that the result
	// is a single file, with no additional .wal or .shm file.
	_, err = tdb.Exec("pragma wal_checkpoint(TRUNCATE)")
	if err != nil {
		sendAPIError(serverError, err.Error(), w, r)
		return err
	}
	tdb.Close()

	// Stream the result to the client
	return streamFile(dbFileName, w)
}

func writeMetadata(pgTx *sql.Tx, tdb *sql.DB, tables map[string]*pgTable) error {
	_, err := tdb.Exec(`
		create table _transicator_metadata
		(key varchar primary key,
		 value varchar)
	 `)
	if err == nil {
		_, err = tdb.Exec(`
			create table _transicator_tables
			(tableName varchar not null,
			 columnName varchar not null,
			 typid integer,
			 primaryKey bool)
		`)
	}

	if err == nil {
		row := pgTx.QueryRow("select txid_current_snapshot()")
		var snap string
		err = row.Scan(&snap)
		if err == nil {
			_, err = tdb.Exec(`
			insert into _transicator_metadata (key, value) values('snapshot', ?)
			`, snap)
		}
	}

	var st *sql.Stmt
	if err == nil {
		st, err = tdb.Prepare(`
			insert into _transicator_tables
			(tableName, columnName, typid, primaryKey)
			values (?, ?, ?, ?)
		`)
		defer st.Close()
	}

	if err == nil {
		for _, table := range tables {
			for _, col := range table.columns {
				_, err = st.Exec(table.schema+"_"+table.name, col.name, col.typid, col.primaryKey)
				if err != nil {
					return err
				}
			}
		}
	}

	return err
}

func createDatabase(fileName string) (*sql.DB, error) {
	log.Debugf("Opening temporary SQLite database in %s", fileName)
	tdb, err := sql.Open("sqlite3", fileName)
	if err == nil {
		err = tdb.Ping()
	}
	if err != nil {
		return nil, err
	}

	_, err = tdb.Exec("pragma journal_mode = WAL")
	if err != nil {
		tdb.Close()
		return nil, err
	}
	return tdb, err
}

func makeSqliteTable(tdb *sql.DB, t *pgTable) error {
	sql := makeSqliteTableSQL(t)
	log.Debugf("New SQLite table: %s", sql)
	_, err := tdb.Exec(sql)
	return err
}

// makeSqliteTableSQL turns the parsed table description from Postgres
// into a "create table" statement that works in SQLite.
func makeSqliteTableSQL(t *pgTable) string {
	s := &bytes.Buffer{}

	s.WriteString(fmt.Sprintf("create table %s (", t.schema+"_"+t.name))
	first := true

	for _, col := range t.columns {
		if first {
			first = false
		} else {
			s.WriteString(",")
		}
		s.WriteString(fmt.Sprintf("%s %s", col.name, convertPgType(col.typid)))
	}
	if len(t.primaryKeys) > 0 {
		s.WriteString(", primary key (")
		for pki, pk := range t.primaryKeys {
			if pki > 0 {
				s.WriteString(",")
			}
			s.WriteString(pk)
		}
		s.WriteString(")")
	}
	s.WriteString(")")
	return s.String()
}

func copyData(pgTx *sql.Tx, tdb *sql.DB, scopes []string, pgTable *pgTable) error {

	sql := fmt.Sprintf("select * from %s.%s where %s in %s",
		pgTable.schema, pgTable.name, selectorColumn, GetTenants(scopes))
	log.Debugf("Postgres query: %s", sql)

	pgRows, err := pgTx.Query(sql)
	if err != nil {
		return err
	}
	defer pgRows.Close()

	// Use column names from the query and not from the table definition.
	// That way we are sure that we are getting them in the right order.
	colNames, colTypes, err := parseColumnNames(pgRows)
	if err != nil {
		return err
	}

	sql = makeInsertSQL(pgTable, colNames)
	log.Debugf("Sqlite insert: %s", sql)

	stmt, err := tdb.Prepare(sql)
	if err != nil {
		return err
	}
	defer stmt.Close()

	for pgRows.Next() {
		cols := make([]interface{}, len(colNames))
		for i := range cols {
			switch pgToSqliteType(int(colTypes[i])) {
			case sqlInteger:
				cols[i] = new(*int64)
			case sqlReal:
				cols[i] = new(*float64)
			case sqlBlob:
				cols[i] = new([]byte)
			case sqlText:
				cols[i] = new(*string)
			case sqlTimestamp:
				cols[i] = new(*time.Time)
			default:
				panic(fmt.Sprintf("Invalid type code %d", pgToSqliteType(int(colTypes[i]))))
			}
		}
		err = pgRows.Scan(cols...)
		if err != nil {
			log.Errorf("Postgres scan error %s. cols = %s", err, debugColTypes(cols))
			return err
		}

		patchColTypes(cols)

		_, err := stmt.Exec(cols...)
		if err != nil {
			log.Errorf("SQLite insert error %s. SQL = %s. cols = %s", err, sql, debugColTypes(cols))
			return err
		}
	}

	return nil
}

func debugColTypes(cols []interface{}) string {
	s := &bytes.Buffer{}
	for i, c := range cols {
		if i > 0 {
			s.WriteString(",")
		}
		s.WriteString(fmt.Sprintf("%T", c))
	}
	return s.String()
}

// patchColTypes looks for column types that are different between PG
// and SQLite.
func patchColTypes(cols []interface{}) {
	for i, c := range cols {
		switch c.(type) {
		case **time.Time:
			ts := *(c.(**time.Time))
			if ts == nil {
				cols[i] = ""
			} else {
				cols[i] = ts.UTC().Format(sqliteTimestampFormat)
			}
		}
	}
}

func makeInsertSQL(pgTable *pgTable, colNames []string) string {
	s := &bytes.Buffer{}
	s.WriteString(fmt.Sprintf("insert into %s (", pgTable.schema+"_"+pgTable.name))

	for i, cn := range colNames {
		if i > 0 {
			s.WriteString(",")
		}
		s.WriteString(cn)
	}

	s.WriteString(") values(")

	for i := range colNames {
		if i > 0 {
			s.WriteString(",")
		}
		s.WriteString("?")
	}
	s.WriteString(")")

	return s.String()
}

func getSHA256Checksum(srcFile string) (error, string) {
	inFileHash, err := os.Open(srcFile)
	if err != nil {
		return err, ""
	}
	defer inFileHash.Close()
	hasher := sha256.New()
	_, err = io.Copy(hasher, inFileHash)
	if err != nil {
		return err, ""
	}
	return nil, hex.EncodeToString(hasher.Sum(nil))
}

func streamFile(srcFile string, w http.ResponseWriter) error {

	err, shasum := getSHA256Checksum(srcFile)
	if err != nil {
		return err
	}
	w.Header().Set("SHA256Sum", shasum)
	w.Header().Set("Content-Type", sqlMediaType)

	inFile, err := os.Open(srcFile)
	if err != nil {
		return err
	}
	defer inFile.Close()
	_, err = io.Copy(w, inFile)

	return err
}
<commit_msg>remove unwanted space.<commit_after>/*
Copyright 2016 The Transicator Authors

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

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

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

import (
	"bytes"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"

	"time"

	log "github.com/Sirupsen/logrus"
	sqlite "github.com/mattn/go-sqlite3"
)

var sqliteTimestampFormat = sqlite.SQLiteTimestampFormats[0]

/*
WriteSqliteSnapshot is responsible for generating the SQLite format of
the snapshot, and also for streaming the file back to the
caller.
*/
func WriteSqliteSnapshot(scopes []string, db *sql.DB, w http.ResponseWriter, r *http.Request) error {

	// Create temporary directory for the temporary DB
	dirName, err := ioutil.TempDir(tempSnapshotDir, tempSnapshotPrefix)
	if err != nil {
		sendAPIError(serverError, err.Error(), w, r)
		return err
	}
	defer func() {
		os.RemoveAll(dirName)
	}()

	// Open and verify the DB
	dbFileName := path.Join(dirName, tempSnapshotName)
	tdb, err := createDatabase(dbFileName)
	if err != nil {
		sendAPIError(http.StatusInternalServerError, err.Error(), w, r)
		return err
	}
	defer tdb.Close()

	// Because of previous config, this puts us in "repeatable read" mode
	pgTx, err := db.Begin()
	if err != nil {
		sendAPIError(serverError, err.Error(), w, r)
		return err
	}
	defer pgTx.Commit()

	tables, err := enumeratePgTables(pgTx)
	if err != nil {
		sendAPIError(http.StatusInternalServerError, err.Error(), w, r)
		return err
	}

	err = writeMetadata(pgTx, tdb, tables)
	if err != nil {
		sendAPIError(http.StatusInternalServerError, err.Error(), w, r)
		return err
	}
	// put tx id into header
	row := pgTx.QueryRow("select txid_current_snapshot()")
	var txId string
	err = row.Scan(&txId)
	if err == nil {
		w.Header().Set("Transicator-Snapshot-TXID", txId)
	}

	// For each table, update the DB
	for tid, pgTable := range tables {
		if pgTable.hasSelector {
			err = makeSqliteTable(tdb, pgTable)
			if err == nil {
				err = copyData(pgTx, tdb, scopes, pgTable)
			}
			if err != nil {
				sendAPIError(serverError, err.Error(), w, r)
				return err
			}
		} else {
			log.Debugf("Skipping table %s which has no selector", tid)
		}
	}

	pgTx.Commit()

	// Close the database. Checkpoint the WAL so that the result
	// is a single file, with no additional .wal or .shm file.
	_, err = tdb.Exec("pragma wal_checkpoint(TRUNCATE)")
	if err != nil {
		sendAPIError(serverError, err.Error(), w, r)
		return err
	}
	tdb.Close()

	// Stream the result to the client
	return streamFile(dbFileName, w)
}

func writeMetadata(pgTx *sql.Tx, tdb *sql.DB, tables map[string]*pgTable) error {
	_, err := tdb.Exec(`
		create table _transicator_metadata
		(key varchar primary key,
		 value varchar)
	 `)
	if err == nil {
		_, err = tdb.Exec(`
			create table _transicator_tables
			(tableName varchar not null,
			 columnName varchar not null,
			 typid integer,
			 primaryKey bool)
		`)
	}

	if err == nil {
		row := pgTx.QueryRow("select txid_current_snapshot()")
		var snap string
		err = row.Scan(&snap)
		if err == nil {
			_, err = tdb.Exec(`
			insert into _transicator_metadata (key, value) values('snapshot', ?)
			`, snap)
		}
	}

	var st *sql.Stmt
	if err == nil {
		st, err = tdb.Prepare(`
			insert into _transicator_tables
			(tableName, columnName, typid, primaryKey)
			values (?, ?, ?, ?)
		`)
		defer st.Close()
	}

	if err == nil {
		for _, table := range tables {
			for _, col := range table.columns {
				_, err = st.Exec(table.schema+"_"+table.name, col.name, col.typid, col.primaryKey)
				if err != nil {
					return err
				}
			}
		}
	}

	return err
}

func createDatabase(fileName string) (*sql.DB, error) {
	log.Debugf("Opening temporary SQLite database in %s", fileName)
	tdb, err := sql.Open("sqlite3", fileName)
	if err == nil {
		err = tdb.Ping()
	}
	if err != nil {
		return nil, err
	}

	_, err = tdb.Exec("pragma journal_mode = WAL")
	if err != nil {
		tdb.Close()
		return nil, err
	}
	return tdb, err
}

func makeSqliteTable(tdb *sql.DB, t *pgTable) error {
	sql := makeSqliteTableSQL(t)
	log.Debugf("New SQLite table: %s", sql)
	_, err := tdb.Exec(sql)
	return err
}

// makeSqliteTableSQL turns the parsed table description from Postgres
// into a "create table" statement that works in SQLite.
func makeSqliteTableSQL(t *pgTable) string {
	s := &bytes.Buffer{}

	s.WriteString(fmt.Sprintf("create table %s (", t.schema+"_"+t.name))
	first := true

	for _, col := range t.columns {
		if first {
			first = false
		} else {
			s.WriteString(",")
		}
		s.WriteString(fmt.Sprintf("%s %s", col.name, convertPgType(col.typid)))
	}
	if len(t.primaryKeys) > 0 {
		s.WriteString(", primary key (")
		for pki, pk := range t.primaryKeys {
			if pki > 0 {
				s.WriteString(",")
			}
			s.WriteString(pk)
		}
		s.WriteString(")")
	}
	s.WriteString(")")
	return s.String()
}

func copyData(pgTx *sql.Tx, tdb *sql.DB, scopes []string, pgTable *pgTable) error {

	sql := fmt.Sprintf("select * from %s.%s where %s in %s",
		pgTable.schema, pgTable.name, selectorColumn, GetTenants(scopes))
	log.Debugf("Postgres query: %s", sql)

	pgRows, err := pgTx.Query(sql)
	if err != nil {
		return err
	}
	defer pgRows.Close()

	// Use column names from the query and not from the table definition.
	// That way we are sure that we are getting them in the right order.
	colNames, colTypes, err := parseColumnNames(pgRows)
	if err != nil {
		return err
	}

	sql = makeInsertSQL(pgTable, colNames)
	log.Debugf("Sqlite insert: %s", sql)

	stmt, err := tdb.Prepare(sql)
	if err != nil {
		return err
	}
	defer stmt.Close()

	for pgRows.Next() {
		cols := make([]interface{}, len(colNames))
		for i := range cols {
			switch pgToSqliteType(int(colTypes[i])) {
			case sqlInteger:
				cols[i] = new(*int64)
			case sqlReal:
				cols[i] = new(*float64)
			case sqlBlob:
				cols[i] = new([]byte)
			case sqlText:
				cols[i] = new(*string)
			case sqlTimestamp:
				cols[i] = new(*time.Time)
			default:
				panic(fmt.Sprintf("Invalid type code %d", pgToSqliteType(int(colTypes[i]))))
			}
		}
		err = pgRows.Scan(cols...)
		if err != nil {
			log.Errorf("Postgres scan error %s. cols = %s", err, debugColTypes(cols))
			return err
		}

		patchColTypes(cols)

		_, err := stmt.Exec(cols...)
		if err != nil {
			log.Errorf("SQLite insert error %s. SQL = %s. cols = %s", err, sql, debugColTypes(cols))
			return err
		}
	}

	return nil
}

func debugColTypes(cols []interface{}) string {
	s := &bytes.Buffer{}
	for i, c := range cols {
		if i > 0 {
			s.WriteString(",")
		}
		s.WriteString(fmt.Sprintf("%T", c))
	}
	return s.String()
}

// patchColTypes looks for column types that are different between PG
// and SQLite.
func patchColTypes(cols []interface{}) {
	for i, c := range cols {
		switch c.(type) {
		case **time.Time:
			ts := *(c.(**time.Time))
			if ts == nil {
				cols[i] = ""
			} else {
				cols[i] = ts.UTC().Format(sqliteTimestampFormat)
			}
		}
	}
}

func makeInsertSQL(pgTable *pgTable, colNames []string) string {
	s := &bytes.Buffer{}
	s.WriteString(fmt.Sprintf("insert into %s (", pgTable.schema+"_"+pgTable.name))

	for i, cn := range colNames {
		if i > 0 {
			s.WriteString(",")
		}
		s.WriteString(cn)
	}

	s.WriteString(") values(")

	for i := range colNames {
		if i > 0 {
			s.WriteString(",")
		}
		s.WriteString("?")
	}
	s.WriteString(")")

	return s.String()
}

func getSHA256Checksum(srcFile string) (error, string) {
	inFileHash, err := os.Open(srcFile)
	if err != nil {
		return err, ""
	}
	defer inFileHash.Close()
	hasher := sha256.New()
	_, err = io.Copy(hasher, inFileHash)
	if err != nil {
		return err, ""
	}
	return nil, hex.EncodeToString(hasher.Sum(nil))
}

func streamFile(srcFile string, w http.ResponseWriter) error {

	err, shasum := getSHA256Checksum(srcFile)
	if err != nil {
		return err
	}
	w.Header().Set("SHA256Sum", shasum)
	w.Header().Set("Content-Type", sqlMediaType)

	inFile, err := os.Open(srcFile)
	if err != nil {
		return err
	}
	defer inFile.Close()
	_, err = io.Copy(w, inFile)
	return err
}
<|endoftext|>


<commit_before>package cli

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"github.com/99designs/aws-vault/server"
	"github.com/99designs/aws-vault/vault"
	"github.com/99designs/keyring"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"gopkg.in/alecthomas/kingpin.v2"
)

type ExecCommandInput struct {
	ProfileName      string
	Command          string
	Args             []string
	Keyring          keyring.Keyring
	StartServer      bool
	CredentialHelper bool
	Config           vault.Config
	SessionDuration  time.Duration
	NoSession        bool
}

// AwsCredentialHelperData is metadata for AWS CLI credential process
// See https://docs.aws.amazon.com/cli/latest/topic/config-vars.html#sourcing-credentials-from-external-processes
type AwsCredentialHelperData struct {
	Version         int    `json:"Version"`
	AccessKeyID     string `json:"AccessKeyId"`
	SecretAccessKey string `json:"SecretAccessKey"`
	SessionToken    string `json:"SessionToken"`
	Expiration      string `json:"Expiration,omitempty"`
}

func ConfigureExecCommand(app *kingpin.Application) {
	input := ExecCommandInput{}

	cmd := app.Command("exec", "Executes a command with AWS credentials in the environment")

	cmd.Flag("duration", "Duration of the temporary or assume-role session. Defaults to 1h").
		Short('d').
		DurationVar(&input.SessionDuration)

	cmd.Flag("no-session", "Use master credentials, no session created").
		Short('n').
		BoolVar(&input.NoSession)

	cmd.Flag("mfa-token", "The MFA token to use").
		Short('t').
		StringVar(&input.Config.MfaToken)

	cmd.Flag("json", "AWS credential helper. Ref: https://docs.aws.amazon.com/cli/latest/topic/config-vars.html#sourcing-credentials-from-external-processes").
		Short('j').
		BoolVar(&input.CredentialHelper)

	cmd.Flag("server", "Run the server in the background for credentials").
		Short('s').
		BoolVar(&input.StartServer)

	cmd.Arg("profile", "Name of the profile").
		Required().
		HintAction(awsConfigFile.ProfileNames).
		StringVar(&input.ProfileName)

	cmd.Arg("cmd", "Command to execute, defaults to $SHELL").
		Default(os.Getenv("SHELL")).
		StringVar(&input.Command)

	cmd.Arg("args", "Command arguments").
		StringsVar(&input.Args)

	cmd.Action(func(c *kingpin.ParseContext) error {
		input.Keyring = keyringImpl
		input.Config.MfaPromptMethod = GlobalFlags.PromptDriver
		input.Config.GetSessionTokenDuration = input.SessionDuration
		input.Config.AssumeRoleDuration = input.SessionDuration
		app.FatalIfError(ExecCommand(input), "exec")
		return nil
	})
}

func ExecCommand(input ExecCommandInput) error {
	if os.Getenv("AWS_VAULT") != "" {
		return fmt.Errorf("aws-vault sessions should be nested with care, unset $AWS_VAULT to force")
	}

	var setEnv = true

	if input.NoSession && input.StartServer {
		return fmt.Errorf("Can't start a credential server without a session")
	}

	configLoader.BaseConfig = input.Config
	configLoader.ProfileNameForEnv = input.ProfileName
	config, err := configLoader.LoadFromProfile(input.ProfileName)
	if err != nil {
		return err
	}

	credKeyring := &vault.CredentialKeyring{Keyring: input.Keyring}
	var creds *credentials.Credentials
	if input.NoSession {
		creds = vault.NewMasterCredentials(credKeyring, config.ProfileName)
	} else {
		creds, err = vault.NewTempCredentials(input.ProfileName, credKeyring, configLoader)
		if err != nil {
			return fmt.Errorf("Error getting temporary credentials: %w", err)
		}
	}

	val, err := creds.Get()
	if err != nil {
		return fmt.Errorf("Failed to get credentials for %s: %w", input.ProfileName, err)
	}

	if input.StartServer {
		if err := server.StartCredentialsServer(creds); err != nil {
			return fmt.Errorf("Failed to start credential server: %w", err)
		} else {
			setEnv = false
		}
	}

	if input.CredentialHelper {
		credentialData := AwsCredentialHelperData{
			Version:         1,
			AccessKeyID:     val.AccessKeyID,
			SecretAccessKey: val.SecretAccessKey,
			SessionToken:    val.SessionToken,
		}
		if !input.NoSession {
			credsExprest, err := creds.ExpiresAt()
			if err != nil {
				return fmt.Errorf("Error getting credential expiration: %w", err)
			}
			credentialData.Expiration = credsExprest.Format("2006-01-02T15:04:05Z")
		}
		json, err := json.Marshal(&credentialData)
		if err != nil {
			return fmt.Errorf("Error creating credential json: %w", err)
		}
		fmt.Print(string(json))
	} else {

		env := environ(os.Environ())
		env.Set("AWS_VAULT", input.ProfileName)

		env.Unset("AWS_ACCESS_KEY_ID")
		env.Unset("AWS_SECRET_ACCESS_KEY")
		env.Unset("AWS_CREDENTIAL_FILE")
		env.Unset("AWS_DEFAULT_PROFILE")
		env.Unset("AWS_PROFILE")

		if config.Region != "" {
			log.Printf("Setting subprocess env: AWS_DEFAULT_REGION=%s, AWS_REGION=%s", config.Region, config.Region)
			env.Set("AWS_DEFAULT_REGION", config.Region)
			env.Set("AWS_REGION", config.Region)
		}

		if setEnv {
			log.Println("Setting subprocess env: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY")
			env.Set("AWS_ACCESS_KEY_ID", val.AccessKeyID)
			env.Set("AWS_SECRET_ACCESS_KEY", val.SecretAccessKey)

			if val.SessionToken != "" {
				log.Println("Setting subprocess env: AWS_SESSION_TOKEN, AWS_SECURITY_TOKEN")
				env.Set("AWS_SESSION_TOKEN", val.SessionToken)
				env.Set("AWS_SECURITY_TOKEN", val.SessionToken)
			}
		}

		err = exec(input.Command, input.Args, env)
		if err != nil {
			return fmt.Errorf("Error execing process: %w", err)
		}
	}

	return nil
}

// environ is a slice of strings representing the environment, in the form "key=value".
type environ []string

// Unset an environment variable by key
func (e *environ) Unset(key string) {
	for i := range *e {
		if strings.HasPrefix((*e)[i], key+"=") {
			(*e)[i] = (*e)[len(*e)-1]
			*e = (*e)[:len(*e)-1]
			break
		}
	}
}

// Set adds an environment variable, replacing any existing ones of the same key
func (e *environ) Set(key, val string) {
	e.Unset(key)
	*e = append(*e, key+"="+val)
}
<commit_msg>Set session expiration env var on exec Fixes #473<commit_after>package cli

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"github.com/99designs/aws-vault/server"
	"github.com/99designs/aws-vault/vault"
	"github.com/99designs/keyring"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"gopkg.in/alecthomas/kingpin.v2"
)

type ExecCommandInput struct {
	ProfileName      string
	Command          string
	Args             []string
	Keyring          keyring.Keyring
	StartServer      bool
	CredentialHelper bool
	Config           vault.Config
	SessionDuration  time.Duration
	NoSession        bool
}

// AwsCredentialHelperData is metadata for AWS CLI credential process
// See https://docs.aws.amazon.com/cli/latest/topic/config-vars.html#sourcing-credentials-from-external-processes
type AwsCredentialHelperData struct {
	Version         int    `json:"Version"`
	AccessKeyID     string `json:"AccessKeyId"`
	SecretAccessKey string `json:"SecretAccessKey"`
	SessionToken    string `json:"SessionToken"`
	Expiration      string `json:"Expiration,omitempty"`
}

func ConfigureExecCommand(app *kingpin.Application) {
	input := ExecCommandInput{}

	cmd := app.Command("exec", "Executes a command with AWS credentials in the environment")

	cmd.Flag("duration", "Duration of the temporary or assume-role session. Defaults to 1h").
		Short('d').
		DurationVar(&input.SessionDuration)

	cmd.Flag("no-session", "Use master credentials, no session created").
		Short('n').
		BoolVar(&input.NoSession)

	cmd.Flag("mfa-token", "The MFA token to use").
		Short('t').
		StringVar(&input.Config.MfaToken)

	cmd.Flag("json", "AWS credential helper. Ref: https://docs.aws.amazon.com/cli/latest/topic/config-vars.html#sourcing-credentials-from-external-processes").
		Short('j').
		BoolVar(&input.CredentialHelper)

	cmd.Flag("server", "Run the server in the background for credentials").
		Short('s').
		BoolVar(&input.StartServer)

	cmd.Arg("profile", "Name of the profile").
		Required().
		HintAction(awsConfigFile.ProfileNames).
		StringVar(&input.ProfileName)

	cmd.Arg("cmd", "Command to execute, defaults to $SHELL").
		Default(os.Getenv("SHELL")).
		StringVar(&input.Command)

	cmd.Arg("args", "Command arguments").
		StringsVar(&input.Args)

	cmd.Action(func(c *kingpin.ParseContext) error {
		input.Keyring = keyringImpl
		input.Config.MfaPromptMethod = GlobalFlags.PromptDriver
		input.Config.GetSessionTokenDuration = input.SessionDuration
		input.Config.AssumeRoleDuration = input.SessionDuration
		app.FatalIfError(ExecCommand(input), "exec")
		return nil
	})
}

func ExecCommand(input ExecCommandInput) error {
	if os.Getenv("AWS_VAULT") != "" {
		return fmt.Errorf("aws-vault sessions should be nested with care, unset $AWS_VAULT to force")
	}

	var setEnv = true

	if input.NoSession && input.StartServer {
		return fmt.Errorf("Can't start a credential server without a session")
	}

	configLoader.BaseConfig = input.Config
	configLoader.ProfileNameForEnv = input.ProfileName
	config, err := configLoader.LoadFromProfile(input.ProfileName)
	if err != nil {
		return err
	}

	credKeyring := &vault.CredentialKeyring{Keyring: input.Keyring}
	var creds *credentials.Credentials
	if input.NoSession {
		creds = vault.NewMasterCredentials(credKeyring, config.ProfileName)
	} else {
		creds, err = vault.NewTempCredentials(input.ProfileName, credKeyring, configLoader)
		if err != nil {
			return fmt.Errorf("Error getting temporary credentials: %w", err)
		}
	}

	val, err := creds.Get()
	if err != nil {
		return fmt.Errorf("Failed to get credentials for %s: %w", input.ProfileName, err)
	}

	if input.StartServer {
		if err := server.StartCredentialsServer(creds); err != nil {
			return fmt.Errorf("Failed to start credential server: %w", err)
		} else {
			setEnv = false
		}
	}

	if input.CredentialHelper {
		credentialData := AwsCredentialHelperData{
			Version:         1,
			AccessKeyID:     val.AccessKeyID,
			SecretAccessKey: val.SecretAccessKey,
			SessionToken:    val.SessionToken,
		}
		if !input.NoSession {
			credsExprest, err := creds.ExpiresAt()
			if err != nil {
				return fmt.Errorf("Error getting credential expiration: %w", err)
			}
			credentialData.Expiration = credsExprest.Format("2006-01-02T15:04:05Z")
		}
		json, err := json.Marshal(&credentialData)
		if err != nil {
			return fmt.Errorf("Error creating credential json: %w", err)
		}
		fmt.Print(string(json))
	} else {

		env := environ(os.Environ())
		env.Set("AWS_VAULT", input.ProfileName)

		env.Unset("AWS_ACCESS_KEY_ID")
		env.Unset("AWS_SECRET_ACCESS_KEY")
		env.Unset("AWS_CREDENTIAL_FILE")
		env.Unset("AWS_DEFAULT_PROFILE")
		env.Unset("AWS_PROFILE")

		if config.Region != "" {
			log.Printf("Setting subprocess env: AWS_DEFAULT_REGION=%s, AWS_REGION=%s", config.Region, config.Region)
			env.Set("AWS_DEFAULT_REGION", config.Region)
			env.Set("AWS_REGION", config.Region)
		}

		if setEnv {
			log.Println("Setting subprocess env: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY")
			env.Set("AWS_ACCESS_KEY_ID", val.AccessKeyID)
			env.Set("AWS_SECRET_ACCESS_KEY", val.SecretAccessKey)

			if val.SessionToken != "" {
				log.Println("Setting subprocess env: AWS_SESSION_TOKEN, AWS_SECURITY_TOKEN")
				env.Set("AWS_SESSION_TOKEN", val.SessionToken)
				env.Set("AWS_SECURITY_TOKEN", val.SessionToken)
				expiration, err := creds.ExpiresAt()
				if err == nil {
					log.Println("Setting subprocess env: AWS_SESSION_EXPIRATION")
					env.Set("AWS_SESSION_EXPIRATION", expiration.Format(time.RFC3339))
				}
			}
		}

		err = exec(input.Command, input.Args, env)
		if err != nil {
			return fmt.Errorf("Error execing process: %w", err)
		}
	}

	return nil
}

// environ is a slice of strings representing the environment, in the form "key=value".
type environ []string

// Unset an environment variable by key
func (e *environ) Unset(key string) {
	for i := range *e {
		if strings.HasPrefix((*e)[i], key+"=") {
			(*e)[i] = (*e)[len(*e)-1]
			*e = (*e)[:len(*e)-1]
			break
		}
	}
}

// Set adds an environment variable, replacing any existing ones of the same key
func (e *environ) Set(key, val string) {
	e.Unset(key)
	*e = append(*e, key+"="+val)
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"os"
	"path"
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/bitrise-io/go-utils/cmdex"
	"github.com/bitrise-io/go-utils/colorstring"
	"github.com/bitrise-io/go-utils/pathutil"
	"github.com/bitrise-io/goinp/goinp"
	"github.com/bitrise-tools/codesigndoc/osxkeychain"
	"github.com/bitrise-tools/codesigndoc/provprofile"
	"github.com/bitrise-tools/codesigndoc/utils"
	"github.com/bitrise-tools/codesigndoc/xcode"
	"github.com/codegangsta/cli"
)

const (
	confExportOutputDirPath = "./codesigndoc_exports"
)

func printFinished() {
	fmt.Println()
	fmt.Println(colorstring.Green("That's all."))
	fmt.Println("You just have to upload the found code signing files and you'll be good to go!")
}

func scan(c *cli.Context) {
	projectPath := c.String(FileParamKey)
	if projectPath == "" {
		askText := `Please drag-and-drop your Xcode Project (` + colorstring.Green(".xcodeproj") + `)
   or Workspace (` + colorstring.Green(".xcworkspace") + `) file, the one you usually open in Xcode,
   then hit Enter.

  (Note: if you have a Workspace file you should most likely use that)`
		fmt.Println()
		projpth, err := goinp.AskForString(askText)
		if err != nil {
			log.Fatalf("Failed to read input: %s", err)
		}
		projectPath = projpth
	}
	log.Debugf("projectPath: %s", projectPath)
	xcodeCmd := xcode.CommandModel{
		ProjectFilePath: projectPath,
	}

	schemeToUse := c.String(SchemeParamKey)
	if schemeToUse == "" {
		log.Println("🔦  Scanning Schemes ...")
		schemes, err := xcodeCmd.ScanSchemes()
		if err != nil {
			log.Fatalf("Failed to scan Schemes: %s", err)
		}
		log.Debugf("schemes: %v", schemes)

		fmt.Println()
		selectedScheme, err := goinp.SelectFromStrings("Select the Scheme you usually use in Xcode", schemes)
		if err != nil {
			log.Fatalf("Failed to select Scheme: %s", err)
		}
		log.Debugf("selected scheme: %v", selectedScheme)
		schemeToUse = selectedScheme
	}
	xcodeCmd.Scheme = schemeToUse

	fmt.Println()
	log.Println("🔦  Running an Xcode Archive, to get all the required code signing settings...")
	codeSigningSettings, err := xcodeCmd.ScanCodeSigningSettings()
	if err != nil {
		log.Fatalf("Failed to detect code signing settings: %s", err)
	}
	log.Debugf("codeSigningSettings: %#v", codeSigningSettings)

	fmt.Println()
	utils.Printlnf("=== Required Identities/Certificates (%d) ===", len(codeSigningSettings.Identities))
	for idx, anIdentity := range codeSigningSettings.Identities {
		utils.Printlnf(" * (%d): %s", idx+1, anIdentity.Title)
	}
	fmt.Println("========================================")

	fmt.Println()
	utils.Printlnf("=== Required Provisioning Profiles (%d) ===", len(codeSigningSettings.ProvProfiles))
	for idx, aProvProfile := range codeSigningSettings.ProvProfiles {
		utils.Printlnf(" * (%d): %s (UUID: %s)", idx+1, aProvProfile.Title, aProvProfile.UUID)
	}
	fmt.Println("======================================")

	//
	// --- Code Signing issue checks / report
	//

	if len(codeSigningSettings.Identities) < 1 {
		log.Fatal("No Code Signing Identity detected!")
	}
	if len(codeSigningSettings.Identities) > 1 {
		log.Warning("More than one Code Signing Identity (certificate) is required to sign your app!")
		log.Warning("You should check your settings and make sure a single Identity/Certificate can be used")
		log.Warning(" for Archiving your app!")
	}

	if len(codeSigningSettings.ProvProfiles) < 1 {
		log.Fatal("No Provisioning Profiles detected!")
	}

	//
	// --- Export
	//

	if !c.Bool(AllowExportParamKey) {
		isShouldExport, err := goinp.AskForBool("Do you want to export these files?")
		if err != nil {
			log.Fatalf("Failed to process your input: %s", err)
		}
		if !isShouldExport {
			printFinished()
			return
		}
	} else {
		log.Debug("Allow Export flag was set - doing export without asking")
	}

	fmt.Println()
	log.Println("Exporting the required Identities (Certificates) ...")
	fmt.Println(" You'll most likely see popups (one for each Identity) from Keychain,")
	fmt.Println(" you will have to accept (Allow) those to be able to export the Identity")
	fmt.Println()

	absExportOutputDirPath, err := pathutil.AbsPath(confExportOutputDirPath)
	log.Debugf("absExportOutputDirPath: %s", absExportOutputDirPath)
	if err != nil {
		log.Fatalf("Failed to determin Absolute path of export dir: %s", confExportOutputDirPath)
	}
	if exist, err := pathutil.IsDirExists(absExportOutputDirPath); err != nil {
		log.Fatalf("Failed to determin whether the export directory already exists: %s", err)
	} else if !exist {
		if err := os.Mkdir(absExportOutputDirPath, 0777); err != nil {
			log.Fatalf("Failed to create export output directory at path: %s | error: %s", absExportOutputDirPath, err)
		}
	} else {
		log.Debugf("Export output dir already exists at path: %s", absExportOutputDirPath)
	}

	identityExportRefs := osxkeychain.CreateEmptyCFTypeRefSlice()
	defer osxkeychain.ReleaseRefList(identityExportRefs)

	fmt.Println()
	for _, aIdentity := range codeSigningSettings.Identities {
		log.Infof(" * Exporting Identity: %s", aIdentity.Title)
		foundIdentityRefs, err := osxkeychain.FindIdentity(aIdentity.Title)
		if err != nil {
			log.Fatalf("Failed to Export Identity: %s", err)
		}
		log.Debugf("identityRefs: %d", len(foundIdentityRefs))
		if len(foundIdentityRefs) < 1 {
			log.Fatalf("No Identity found in Keychain!")
		}

		// check validity
		validIdentityRefs := osxkeychain.CreateEmptyCFTypeRefSlice()
		for _, aIdentityRef := range foundIdentityRefs {
			cert, err := osxkeychain.GetCertificateDataFromIdentityRef(aIdentityRef)
			if err != nil {
				log.Fatalf("Failed to read certificate data: %s", err)
			}
			log.Debugf("cert - valid - NotBefore: %#v", cert.NotBefore)
			log.Debugf("cert - valid - NotAfter: %#v", cert.NotAfter)

			timeNow := time.Now()
			if !timeNow.After(cert.NotBefore) {
				log.Warningf("Certificate is not yet valid - validity starts at: %s", cert.NotBefore)
				continue
			}
			if !timeNow.Before(cert.NotAfter) {
				log.Warningf("Certificate is not valid anymore - validity ended at: %s", cert.NotAfter)
				continue
			}
			log.Debugf("Certificate is Valid, based on it's validity date-times")
			validIdentityRefs = append(validIdentityRefs, aIdentityRef)
		}

		if len(validIdentityRefs) < 1 {
			log.Fatalf("Identity found found in Keychain, but no Valid identity found!")
		}
		if len(validIdentityRefs) > 1 {
			log.Fatalf("Multiple matching Identities found in Keychain! Most likely you have duplicate identity in separate Keychains, like one in System.keychain and one in your Login.keychain")
		}
		identityExportRefs = append(identityExportRefs, validIdentityRefs...)
	}

	if err := osxkeychain.ExportFromKeychain(identityExportRefs, path.Join(absExportOutputDirPath, "Identities.p12")); err != nil {
		log.Fatalf("Failed to export from Keychain: %s", err)
	}

	fmt.Println()
	for _, aProvProfile := range codeSigningSettings.ProvProfiles {
		log.Infof(" * Exporting Provisioning Profile: %s (UUID: %s)", aProvProfile.Title, aProvProfile.UUID)
		filePth, err := provprofile.FindProvProfileFile(aProvProfile)
		if err != nil {
			log.Fatalf("Failed to find Provisioning Profile: %s", err)
		}
		log.Infof("  File found at: %s", filePth)

		if err := cmdex.RunCommand("cp", filePth, absExportOutputDirPath+"/"); err != nil {
			log.Fatalf("Failed to copy the Provisioning Profile into the export directory: %s", err)
		}

		// if err := provprofile.PrintFileInfo(filePth); err != nil {
		// 	log.Fatalf("Err: %s", err)
		// }
	}

	fmt.Println()
	fmt.Printf(colorstring.Green("Exports finished")+" you can find the exported files at: %s\n", absExportOutputDirPath)
	if err := cmdex.RunCommand("open", absExportOutputDirPath); err != nil {
		log.Errorf("Failed to open the export directory in Finder: %s", absExportOutputDirPath)
	}
	fmt.Println("Opened the directory in Finder.")
	fmt.Println()

	printFinished()
}
<commit_msg>typo fix, more log<commit_after>package cli

import (
	"fmt"
	"os"
	"path"
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/bitrise-io/go-utils/cmdex"
	"github.com/bitrise-io/go-utils/colorstring"
	"github.com/bitrise-io/go-utils/pathutil"
	"github.com/bitrise-io/goinp/goinp"
	"github.com/bitrise-tools/codesigndoc/osxkeychain"
	"github.com/bitrise-tools/codesigndoc/provprofile"
	"github.com/bitrise-tools/codesigndoc/utils"
	"github.com/bitrise-tools/codesigndoc/xcode"
	"github.com/codegangsta/cli"
)

const (
	confExportOutputDirPath = "./codesigndoc_exports"
)

func printFinished() {
	fmt.Println()
	fmt.Println(colorstring.Green("That's all."))
	fmt.Println("You just have to upload the found code signing files and you'll be good to go!")
}

func scan(c *cli.Context) {
	projectPath := c.String(FileParamKey)
	if projectPath == "" {
		askText := `Please drag-and-drop your Xcode Project (` + colorstring.Green(".xcodeproj") + `)
   or Workspace (` + colorstring.Green(".xcworkspace") + `) file, the one you usually open in Xcode,
   then hit Enter.

  (Note: if you have a Workspace file you should most likely use that)`
		fmt.Println()
		projpth, err := goinp.AskForString(askText)
		if err != nil {
			log.Fatalf("Failed to read input: %s", err)
		}
		projectPath = projpth
	}
	log.Debugf("projectPath: %s", projectPath)
	xcodeCmd := xcode.CommandModel{
		ProjectFilePath: projectPath,
	}

	schemeToUse := c.String(SchemeParamKey)
	if schemeToUse == "" {
		log.Println("🔦  Scanning Schemes ...")
		schemes, err := xcodeCmd.ScanSchemes()
		if err != nil {
			log.Fatalf("Failed to scan Schemes: %s", err)
		}
		log.Debugf("schemes: %v", schemes)

		fmt.Println()
		selectedScheme, err := goinp.SelectFromStrings("Select the Scheme you usually use in Xcode", schemes)
		if err != nil {
			log.Fatalf("Failed to select Scheme: %s", err)
		}
		log.Debugf("selected scheme: %v", selectedScheme)
		schemeToUse = selectedScheme
	}
	xcodeCmd.Scheme = schemeToUse

	fmt.Println()
	log.Println("🔦  Running an Xcode Archive, to get all the required code signing settings...")
	codeSigningSettings, err := xcodeCmd.ScanCodeSigningSettings()
	if err != nil {
		log.Fatalf("Failed to detect code signing settings: %s", err)
	}
	log.Debugf("codeSigningSettings: %#v", codeSigningSettings)

	fmt.Println()
	utils.Printlnf("=== Required Identities/Certificates (%d) ===", len(codeSigningSettings.Identities))
	for idx, anIdentity := range codeSigningSettings.Identities {
		utils.Printlnf(" * (%d): %s", idx+1, anIdentity.Title)
	}
	fmt.Println("========================================")

	fmt.Println()
	utils.Printlnf("=== Required Provisioning Profiles (%d) ===", len(codeSigningSettings.ProvProfiles))
	for idx, aProvProfile := range codeSigningSettings.ProvProfiles {
		utils.Printlnf(" * (%d): %s (UUID: %s)", idx+1, aProvProfile.Title, aProvProfile.UUID)
	}
	fmt.Println("======================================")

	//
	// --- Code Signing issue checks / report
	//

	if len(codeSigningSettings.Identities) < 1 {
		log.Fatal("No Code Signing Identity detected!")
	}
	if len(codeSigningSettings.Identities) > 1 {
		log.Warning("More than one Code Signing Identity (certificate) is required to sign your app!")
		log.Warning("You should check your settings and make sure a single Identity/Certificate can be used")
		log.Warning(" for Archiving your app!")
	}

	if len(codeSigningSettings.ProvProfiles) < 1 {
		log.Fatal("No Provisioning Profiles detected!")
	}

	//
	// --- Export
	//

	if !c.Bool(AllowExportParamKey) {
		isShouldExport, err := goinp.AskForBool("Do you want to export these files?")
		if err != nil {
			log.Fatalf("Failed to process your input: %s", err)
		}
		if !isShouldExport {
			printFinished()
			return
		}
	} else {
		log.Debug("Allow Export flag was set - doing export without asking")
	}

	fmt.Println()
	log.Println("Exporting the required Identities (Certificates) ...")
	fmt.Println(" You'll most likely see popups (one for each Identity) from Keychain,")
	fmt.Println(" you will have to accept (Allow) those to be able to export the Identity")
	fmt.Println()

	absExportOutputDirPath, err := pathutil.AbsPath(confExportOutputDirPath)
	log.Debugf("absExportOutputDirPath: %s", absExportOutputDirPath)
	if err != nil {
		log.Fatalf("Failed to determin Absolute path of export dir: %s", confExportOutputDirPath)
	}
	if exist, err := pathutil.IsDirExists(absExportOutputDirPath); err != nil {
		log.Fatalf("Failed to determin whether the export directory already exists: %s", err)
	} else if !exist {
		if err := os.Mkdir(absExportOutputDirPath, 0777); err != nil {
			log.Fatalf("Failed to create export output directory at path: %s | error: %s", absExportOutputDirPath, err)
		}
	} else {
		log.Debugf("Export output dir already exists at path: %s", absExportOutputDirPath)
	}

	identityExportRefs := osxkeychain.CreateEmptyCFTypeRefSlice()
	defer osxkeychain.ReleaseRefList(identityExportRefs)

	fmt.Println()
	for _, aIdentity := range codeSigningSettings.Identities {
		log.Infof(" * Exporting Identity: %s", aIdentity.Title)
		foundIdentityRefs, err := osxkeychain.FindIdentity(aIdentity.Title)
		if err != nil {
			log.Fatalf("Failed to Export Identity: %s", err)
		}
		log.Debugf("foundIdentityRefs: %d", len(foundIdentityRefs))
		if len(foundIdentityRefs) < 1 {
			log.Fatalf("No Identity found in Keychain!")
		}

		// check validity
		validIdentityRefs := osxkeychain.CreateEmptyCFTypeRefSlice()
		for _, aIdentityRef := range foundIdentityRefs {
			cert, err := osxkeychain.GetCertificateDataFromIdentityRef(aIdentityRef)
			if err != nil {
				log.Fatalf("Failed to read certificate data: %s", err)
			}
			log.Debugf("cert - valid - NotBefore: %#v", cert.NotBefore)
			log.Debugf("cert - valid - NotAfter: %#v", cert.NotAfter)

			timeNow := time.Now()
			log.Debugf("cert - time now: %#v", timeNow)
			if !timeNow.After(cert.NotBefore) {
				log.Warningf("Certificate is not yet valid - validity starts at: %s", cert.NotBefore)
				continue
			}
			if !timeNow.Before(cert.NotAfter) {
				log.Warningf("Certificate is not valid anymore - validity ended at: %s", cert.NotAfter)
				continue
			}
			log.Debugf("Certificate is Valid, based on it's validity date-times")
			validIdentityRefs = append(validIdentityRefs, aIdentityRef)
		}

		if len(validIdentityRefs) < 1 {
			log.Fatalf("Identity found found in Keychain, but no Valid identity found!")
		}
		if len(validIdentityRefs) > 1 {
			log.Fatalf("Multiple matching Identities found in Keychain! Most likely you have duplicate identity in separate Keychains, like one in System.keychain and one in your Login.keychain")
		}
		identityExportRefs = append(identityExportRefs, validIdentityRefs...)
	}

	if err := osxkeychain.ExportFromKeychain(identityExportRefs, path.Join(absExportOutputDirPath, "Identities.p12")); err != nil {
		log.Fatalf("Failed to export from Keychain: %s", err)
	}

	fmt.Println()
	for _, aProvProfile := range codeSigningSettings.ProvProfiles {
		log.Infof(" * Exporting Provisioning Profile: %s (UUID: %s)", aProvProfile.Title, aProvProfile.UUID)
		filePth, err := provprofile.FindProvProfileFile(aProvProfile)
		if err != nil {
			log.Fatalf("Failed to find Provisioning Profile: %s", err)
		}
		log.Infof("  File found at: %s", filePth)

		if err := cmdex.RunCommand("cp", filePth, absExportOutputDirPath+"/"); err != nil {
			log.Fatalf("Failed to copy the Provisioning Profile into the export directory: %s", err)
		}

		// if err := provprofile.PrintFileInfo(filePth); err != nil {
		// 	log.Fatalf("Err: %s", err)
		// }
	}

	fmt.Println()
	fmt.Printf(colorstring.Green("Exports finished")+" you can find the exported files at: %s\n", absExportOutputDirPath)
	if err := cmdex.RunCommand("open", absExportOutputDirPath); err != nil {
		log.Errorf("Failed to open the export directory in Finder: %s", absExportOutputDirPath)
	}
	fmt.Println("Opened the directory in Finder.")
	fmt.Println()

	printFinished()
}
<|endoftext|>


<commit_before>package capn

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
	"math"
)

var (
	errBufferCall     = errors.New("capn: can't call on a memory buffer")
	ErrInvalidSegment = errors.New("capn: invalid segment id")
	ErrTooMuchData    = errors.New("capn: too much data in stream")
)

type buffer Segment

// NewBuffer creates an expanding single segment buffer. Creating new objects
// will expand the buffer. Data can be nil (or length 0 with some capacity) if
// creating a new session. If parsing an existing segment then data should be
// the segment contents and will not be copied.
func NewBuffer(data []byte) *Segment {
	if uint64(len(data)) > uint64(math.MaxUint32) {
		return nil
	}

	b := &buffer{}
	b.Message = b
	b.Data = data
	return (*Segment)(b)
}

func (b *buffer) NewSegment(minsz int) (*Segment, error) {
	if minsz < 4096 {
		minsz = 4096
	}
	if uint64(len(b.Data)) > uint64(math.MaxUint32)-uint64(minsz) {
		return nil, ErrOverlarge
	}
	b.Data = append(b.Data, make([]byte, minsz)...)
	b.Data = b.Data[:len(b.Data)-minsz]
	return (*Segment)(b), nil
}

func (b *buffer) Lookup(segid uint32) (*Segment, error) {
	if segid == 0 {
		return (*Segment)(b), nil
	} else {
		return nil, ErrInvalidSegment
	}
}

type multiBuffer struct {
	segments []*Segment
}

// NewmultiBuffer creates a new multi segment message. Creating new objects
// will try and reuse the buffers available, but will create new ones if there
// is insufficient capacity. When parsing an existing message data should be
// the list of segments. The data buffers will not be copied.
func NewmultiBuffer(data [][]byte) *Segment {
	m := &multiBuffer{make([]*Segment, len(data))}
	for i, d := range data {
		m.segments[i] = &Segment{m, d, uint32(i), false}
	}
	if len(data) > 0 {
		return m.segments[0]
	}
	return &Segment{Message: m, Data: nil, Id: 0xFFFFFFFF, RootDone: false}
}

var (
	MaxSegmentNumber = 1024
	MaxTotalSize     = 1024 * 1024 * 1024
)

func (m *multiBuffer) NewSegment(minsz int) (*Segment, error) {
	for _, s := range m.segments {
		if len(s.Data)+minsz <= cap(s.Data) {
			return s, nil
		}
	}

	if minsz < 4096 {
		minsz = 4096
	}
	s := &Segment{m, make([]byte, 0, minsz), uint32(len(m.segments)), false}
	m.segments = append(m.segments, s)
	return s, nil
}

func (m *multiBuffer) Lookup(segid uint32) (*Segment, error) {
	if uint(segid) < uint(len(m.segments)) {
		return m.segments[segid], nil
	} else {
		return nil, ErrInvalidSegment
	}
}

// ReadFromStream reads a non-packed serialized stream from r. buf is used to
// buffer the read contents, can be nil, and is provided so that the buffer
// can be reused between messages. The returned segment is the first segment
// read, which contains the root pointer.
func ReadFromStream(r io.Reader, buf *bytes.Buffer) (*Segment, error) {
	if buf == nil {
		buf = new(bytes.Buffer)
	} else {
		buf.Reset()
	}

	if _, err := io.CopyN(buf, r, 4); err != nil {
		return nil, err
	}

	if binary.LittleEndian.Uint32(buf.Bytes()[:]) >= uint32(MaxSegmentNumber) {
		return nil, ErrTooMuchData
	}

	segnum := int(binary.LittleEndian.Uint32(buf.Bytes()[:]) + 1)
	hdrsz := 8*(segnum/2) + 4

	if _, err := io.CopyN(buf, r, int64(hdrsz)); err != nil {
		return nil, err
	}

	total := 0
	for i := 0; i < segnum; i++ {
		sz := binary.LittleEndian.Uint32(buf.Bytes()[4*i+4:])
		if uint64(total)+uint64(sz)*8 > uint64(MaxTotalSize) {
			return nil, ErrTooMuchData
		}
		total += int(sz) * 8
	}

	if _, err := io.CopyN(buf, r, int64(total)); err != nil {
		return nil, err
	}

	hdrv := buf.Bytes()[4 : hdrsz+4]
	datav := buf.Bytes()[hdrsz+4:]

	if segnum == 1 {
		sz := int(binary.LittleEndian.Uint32(hdrv)) * 8
		return NewBuffer(datav[:sz]), nil
	}

	m := &multiBuffer{make([]*Segment, segnum)}
	for i := 0; i < segnum; i++ {
		sz := int(binary.LittleEndian.Uint32(hdrv[4*i:])) * 8
		m.segments[i] = &Segment{m, datav[:sz], uint32(i), false}
		datav = datav[sz:]
	}

	return m.segments[0], nil
}

// ReadFromMemoryZeroCopy: like ReadFromStream, but reads a non-packed
// serialized stream that already resides in memory in the argument data.
// The returned segment is the first segment read, which contains
// the root pointer. The returned bytesRead says how many bytes were
// consumed from data in making seg. The caller should advance the
// data slice by doing data = data[bytesRead:] between successive calls
// to ReadFromMemoryZeroCopy().
func ReadFromMemoryZeroCopy(data []byte) (seg *Segment, bytesRead int64, err error) {

	if len(data) < 4 {
		return nil, 0, io.EOF
	}

	if binary.LittleEndian.Uint32(data[0:4]) >= uint32(MaxSegmentNumber) {
		return nil, 0, ErrTooMuchData
	}

	segnum := int(binary.LittleEndian.Uint32(data[0:4]) + 1)
	hdrsz := 8*(segnum/2) + 4

	b := data[0:(hdrsz + 4)]

	total := 0
	for i := 0; i < segnum; i++ {
		sz := binary.LittleEndian.Uint32(b[4*i+4:])
		if uint64(total)+uint64(sz)*8 > uint64(MaxTotalSize) {
			return nil, 0, ErrTooMuchData
		}
		total += int(sz) * 8
	}
	if total == 0 {
		return nil, 0, io.EOF
	}

	hdrv := data[4:(hdrsz + 4)]
	datav := data[hdrsz+4:]
	m := &multiBuffer{make([]*Segment, segnum)}
	for i := 0; i < segnum; i++ {
		sz := int(binary.LittleEndian.Uint32(hdrv[4*i:])) * 8
		m.segments[i] = &Segment{m, datav[:sz], uint32(i), false}
		datav = datav[sz:]
	}

	return m.segments[0], int64(4 + hdrsz + total), nil
}

// WriteTo writes the message that the segment is part of to the
// provided stream in serialized form.
func (s *Segment) WriteTo(w io.Writer) (int64, error) {
	segnum := uint32(1)
	for {
		if seg, _ := s.Message.Lookup(segnum); seg == nil {
			break
		}
		segnum++
	}

	hdrv := make([]uint8, 8*(segnum/2)+8)
	binary.LittleEndian.PutUint32(hdrv, segnum-1)
	for i := uint32(0); i < segnum; i++ {
		seg, _ := s.Message.Lookup(i)
		binary.LittleEndian.PutUint32(hdrv[4*i+4:], uint32(len(seg.Data)/8))
	}

	if n, err := w.Write(hdrv); err != nil {
		return int64(n), err
	}
	written := int64(len(hdrv))

	for i := uint32(0); i < segnum; i++ {
		seg, _ := s.Message.Lookup(i)
		if n, err := w.Write(seg.Data); err != nil {
			return written + int64(n), err
		} else {
			written += int64(n)
		}
	}

	return written, nil
}
<commit_msg>add warning about buf reuse in ReadFromStream()<commit_after>package capn

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
	"math"
)

var (
	errBufferCall     = errors.New("capn: can't call on a memory buffer")
	ErrInvalidSegment = errors.New("capn: invalid segment id")
	ErrTooMuchData    = errors.New("capn: too much data in stream")
)

type buffer Segment

// NewBuffer creates an expanding single segment buffer. Creating new objects
// will expand the buffer. Data can be nil (or length 0 with some capacity) if
// creating a new session. If parsing an existing segment then data should be
// the segment contents and will not be copied.
func NewBuffer(data []byte) *Segment {
	if uint64(len(data)) > uint64(math.MaxUint32) {
		return nil
	}

	b := &buffer{}
	b.Message = b
	b.Data = data
	return (*Segment)(b)
}

func (b *buffer) NewSegment(minsz int) (*Segment, error) {
	if minsz < 4096 {
		minsz = 4096
	}
	if uint64(len(b.Data)) > uint64(math.MaxUint32)-uint64(minsz) {
		return nil, ErrOverlarge
	}
	b.Data = append(b.Data, make([]byte, minsz)...)
	b.Data = b.Data[:len(b.Data)-minsz]
	return (*Segment)(b), nil
}

func (b *buffer) Lookup(segid uint32) (*Segment, error) {
	if segid == 0 {
		return (*Segment)(b), nil
	} else {
		return nil, ErrInvalidSegment
	}
}

type multiBuffer struct {
	segments []*Segment
}

// NewmultiBuffer creates a new multi segment message. Creating new objects
// will try and reuse the buffers available, but will create new ones if there
// is insufficient capacity. When parsing an existing message data should be
// the list of segments. The data buffers will not be copied.
func NewmultiBuffer(data [][]byte) *Segment {
	m := &multiBuffer{make([]*Segment, len(data))}
	for i, d := range data {
		m.segments[i] = &Segment{m, d, uint32(i), false}
	}
	if len(data) > 0 {
		return m.segments[0]
	}
	return &Segment{Message: m, Data: nil, Id: 0xFFFFFFFF, RootDone: false}
}

var (
	MaxSegmentNumber = 1024
	MaxTotalSize     = 1024 * 1024 * 1024
)

func (m *multiBuffer) NewSegment(minsz int) (*Segment, error) {
	for _, s := range m.segments {
		if len(s.Data)+minsz <= cap(s.Data) {
			return s, nil
		}
	}

	if minsz < 4096 {
		minsz = 4096
	}
	s := &Segment{m, make([]byte, 0, minsz), uint32(len(m.segments)), false}
	m.segments = append(m.segments, s)
	return s, nil
}

func (m *multiBuffer) Lookup(segid uint32) (*Segment, error) {
	if uint(segid) < uint(len(m.segments)) {
		return m.segments[segid], nil
	} else {
		return nil, ErrInvalidSegment
	}
}

// ReadFromStream reads a non-packed serialized stream from r. buf is used to
// buffer the read contents, can be nil, and is provided so that the buffer
// can be reused between messages. The returned segment is the first segment
// read, which contains the root pointer.
//
// Warning about buf reuse:  It is safer to just pass nil for buf.
// When making multiple calls to ReadFromStream() with the same buf argument, you
// may overwrite the data in a previously returned Segment.
// The re-use of buf is an optimization for when you are actually
// done with any previously returned Segment which may have data still alive
// in buf.
//
func ReadFromStream(r io.Reader, buf *bytes.Buffer) (*Segment, error) {
	if buf == nil {
		buf = new(bytes.Buffer)
	} else {
		buf.Reset()
	}

	if _, err := io.CopyN(buf, r, 4); err != nil {
		return nil, err
	}

	if binary.LittleEndian.Uint32(buf.Bytes()[:]) >= uint32(MaxSegmentNumber) {
		return nil, ErrTooMuchData
	}

	segnum := int(binary.LittleEndian.Uint32(buf.Bytes()[:]) + 1)
	hdrsz := 8*(segnum/2) + 4

	if _, err := io.CopyN(buf, r, int64(hdrsz)); err != nil {
		return nil, err
	}

	total := 0
	for i := 0; i < segnum; i++ {
		sz := binary.LittleEndian.Uint32(buf.Bytes()[4*i+4:])
		if uint64(total)+uint64(sz)*8 > uint64(MaxTotalSize) {
			return nil, ErrTooMuchData
		}
		total += int(sz) * 8
	}

	if _, err := io.CopyN(buf, r, int64(total)); err != nil {
		return nil, err
	}

	hdrv := buf.Bytes()[4 : hdrsz+4]
	datav := buf.Bytes()[hdrsz+4:]

	if segnum == 1 {
		sz := int(binary.LittleEndian.Uint32(hdrv)) * 8
		return NewBuffer(datav[:sz]), nil
	}

	m := &multiBuffer{make([]*Segment, segnum)}
	for i := 0; i < segnum; i++ {
		sz := int(binary.LittleEndian.Uint32(hdrv[4*i:])) * 8
		m.segments[i] = &Segment{m, datav[:sz], uint32(i), false}
		datav = datav[sz:]
	}

	return m.segments[0], nil
}

// ReadFromMemoryZeroCopy: like ReadFromStream, but reads a non-packed
// serialized stream that already resides in memory in the argument data.
// The returned segment is the first segment read, which contains
// the root pointer. The returned bytesRead says how many bytes were
// consumed from data in making seg. The caller should advance the
// data slice by doing data = data[bytesRead:] between successive calls
// to ReadFromMemoryZeroCopy().
func ReadFromMemoryZeroCopy(data []byte) (seg *Segment, bytesRead int64, err error) {

	if len(data) < 4 {
		return nil, 0, io.EOF
	}

	if binary.LittleEndian.Uint32(data[0:4]) >= uint32(MaxSegmentNumber) {
		return nil, 0, ErrTooMuchData
	}

	segnum := int(binary.LittleEndian.Uint32(data[0:4]) + 1)
	hdrsz := 8*(segnum/2) + 4

	b := data[0:(hdrsz + 4)]

	total := 0
	for i := 0; i < segnum; i++ {
		sz := binary.LittleEndian.Uint32(b[4*i+4:])
		if uint64(total)+uint64(sz)*8 > uint64(MaxTotalSize) {
			return nil, 0, ErrTooMuchData
		}
		total += int(sz) * 8
	}
	if total == 0 {
		return nil, 0, io.EOF
	}

	hdrv := data[4:(hdrsz + 4)]
	datav := data[hdrsz+4:]
	m := &multiBuffer{make([]*Segment, segnum)}
	for i := 0; i < segnum; i++ {
		sz := int(binary.LittleEndian.Uint32(hdrv[4*i:])) * 8
		m.segments[i] = &Segment{m, datav[:sz], uint32(i), false}
		datav = datav[sz:]
	}

	return m.segments[0], int64(4 + hdrsz + total), nil
}

// WriteTo writes the message that the segment is part of to the
// provided stream in serialized form.
func (s *Segment) WriteTo(w io.Writer) (int64, error) {
	segnum := uint32(1)
	for {
		if seg, _ := s.Message.Lookup(segnum); seg == nil {
			break
		}
		segnum++
	}

	hdrv := make([]uint8, 8*(segnum/2)+8)
	binary.LittleEndian.PutUint32(hdrv, segnum-1)
	for i := uint32(0); i < segnum; i++ {
		seg, _ := s.Message.Lookup(i)
		binary.LittleEndian.PutUint32(hdrv[4*i+4:], uint32(len(seg.Data)/8))
	}

	if n, err := w.Write(hdrv); err != nil {
		return int64(n), err
	}
	written := int64(len(hdrv))

	for i := uint32(0); i < segnum; i++ {
		seg, _ := s.Message.Lookup(i)
		if n, err := w.Write(seg.Data); err != nil {
			return written + int64(n), err
		} else {
			written += int64(n)
		}
	}

	return written, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel

import (
	"errors"
	"sync"

	"github.com/uber/tchannel-go/typed"
	"golang.org/x/net/context"
)

var (
	errDuplicateMex        = errors.New("multiple attempts to use the message id")
	errMexChannelFull      = NewSystemError(ErrCodeBusy, "cannot send frame to message exchange channel")
	errUnexpectedFrameType = errors.New("unexpected frame received")
)

const (
	messageExchangeSetInbound  = "inbound"
	messageExchangeSetOutbound = "outbound"

	// mexChannelBufferSize is the size of the message exchange channel buffer.
	mexChannelBufferSize = 2
)

// A messageExchange tracks this Connections's side of a message exchange with a
// peer.  Each message exchange has a channel that can be used to receive
// frames from the peer, and a Context that can controls when the exchange has
// timed out or been cancelled.
type messageExchange struct {
	recvCh    chan *Frame
	ctx       context.Context
	msgID     uint32
	msgType   messageType
	mexset    *messageExchangeSet
	framePool FramePool
}

// forwardPeerFrame forwards a frame from a peer to the message exchange, where
// it can be pulled by whatever application thread is handling the exchange
func (mex *messageExchange) forwardPeerFrame(frame *Frame) error {
	if err := mex.ctx.Err(); err != nil {
		return GetContextError(err)
	}
	select {
	case mex.recvCh <- frame:
		return nil
	case <-mex.ctx.Done():
		// Note: One slow reader processing a large request could stall the connection.
		// If we see this, we need to increase the recvCh buffer size.
		return GetContextError(mex.ctx.Err())
	}
}

// recvPeerFrame waits for a new frame from the peer, or until the context
// expires or is cancelled
func (mex *messageExchange) recvPeerFrame() (*Frame, error) {
	if err := mex.ctx.Err(); err != nil {
		return nil, GetContextError(err)
	}

	select {
	case frame := <-mex.recvCh:
		if frame.Header.ID != mex.msgID {
			mex.mexset.log.WithFields(
				LogField{"msgId", mex.msgID},
				LogField{"header", frame.Header},
			).Error("recvPeerFrame received msg with unexpected ID.")
			return nil, errUnexpectedFrameType
		}
		return frame, nil
	case <-mex.ctx.Done():
		return nil, GetContextError(mex.ctx.Err())
	}
}

// recvPeerFrameOfType waits for a new frame of a given type from the peer, failing
// if the next frame received is not of that type.
// If an error frame is returned, then the errorMessage is returned as the error.
func (mex *messageExchange) recvPeerFrameOfType(msgType messageType) (*Frame, error) {
	frame, err := mex.recvPeerFrame()
	if err != nil {
		return nil, err
	}

	switch frame.Header.messageType {
	case msgType:
		return frame, nil

	case messageTypeError:
		// If we read an error frame, we can release it once we deserialize it.
		defer mex.framePool.Release(frame)

		errMsg := errorMessage{
			id: frame.Header.ID,
		}
		var rbuf typed.ReadBuffer
		rbuf.Wrap(frame.SizedPayload())
		if err := errMsg.read(&rbuf); err != nil {
			return nil, err
		}
		return nil, errMsg

	default:
		// TODO(mmihic): Should be treated as a protocol error
		mex.mexset.log.WithFields(
			LogField{"header", frame.Header},
			LogField{"expectedType", msgType},
			LogField{"expectedID", mex.msgID},
		).Warn("Received unexpected frame.")
		return nil, errUnexpectedFrameType
	}
}

// shutdown shuts down the message exchange, removing it from the message
// exchange set so  that it cannot receive more messages from the peer.  The
// receive channel remains open, however, in case there are concurrent
// goroutines sending to it.
func (mex *messageExchange) shutdown() {
	mex.mexset.removeExchange(mex.msgID)
}

// inboundTimeout is called when an exchange times out, but a handler may still be
// running in the background. Since the handler may still write to the exchange, we
// cannot shutdown the exchange, but we should remove it from the connection's
// exchange list.
func (mex *messageExchange) inboundTimeout() {
	mex.mexset.timeoutExchange(mex.msgID)
}

// A messageExchangeSet manages a set of active message exchanges.  It is
// mainly used to route frames from a peer to the appropriate messageExchange,
// or to cancel or mark a messageExchange as being in error.  Each Connection
// maintains two messageExchangeSets, one to manage exchanges that it has
// initiated (outbound), and another to manage exchanges that the peer has
// initiated (inbound).  The message-type specific handlers are responsible for
// ensuring that their message exchanges are properly registered and removed
// from the corresponding exchange set.
type messageExchangeSet struct {
	sync.RWMutex

	log        Logger
	name       string
	onRemoved  func()
	onAdded    func()
	sendChRefs sync.WaitGroup

	// maps are mutable, and are protected by the mutex.
	exchanges        map[uint32]*messageExchange
	timeoutExchanges map[uint32]struct{}
}

// newMessageExchangeSet creates a new messageExchangeSet with a given name.
func newMessageExchangeSet(log Logger, name string) *messageExchangeSet {
	return &messageExchangeSet{
		name:             name,
		log:              log.WithFields(LogField{"exchange", name}),
		exchanges:        make(map[uint32]*messageExchange),
		timeoutExchanges: make(map[uint32]struct{}),
	}
}

// newExchange creates and adds a new message exchange to this set
func (mexset *messageExchangeSet) newExchange(ctx context.Context, framePool FramePool,
	msgType messageType, msgID uint32, bufferSize int) (*messageExchange, error) {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("Creating new %s message exchange for [%v:%d]", mexset.name, msgType, msgID)
	}

	mex := &messageExchange{
		msgType:   msgType,
		msgID:     msgID,
		ctx:       ctx,
		recvCh:    make(chan *Frame, bufferSize),
		mexset:    mexset,
		framePool: framePool,
	}

	mexset.Lock()
	if existingMex := mexset.exchanges[mex.msgID]; existingMex != nil {
		if existingMex == mex {
			mexset.log.WithFields(
				LogField{"name", mexset.name},
				LogField{"msgType", mex.msgType},
				LogField{"msgID", mex.msgID},
			).Warn("mex registered multiple times.")
		} else {
			mexset.log.WithFields(
				LogField{"msgID", mex.msgID},
				LogField{"existingType", existingMex.msgType},
				LogField{"newType", mex.msgType},
			).Warn("Duplicate msg ID for active and new mex.")
		}

		mexset.Unlock()
		return nil, errDuplicateMex
	}

	mexset.exchanges[mex.msgID] = mex
	mexset.sendChRefs.Add(1)
	mexset.Unlock()

	mexset.onAdded()

	// TODO(mmihic): Put into a deadline ordered heap so we can garbage collected expired exchanges
	return mex, nil
}

// deleteExchange will delete msgID, and return whether it was found or whether it was
// timed out. This method must be called with the lock.
func (mexset *messageExchangeSet) deleteExchange(msgID uint32) (found, timedOut bool) {
	if _, found := mexset.exchanges[msgID]; found {
		delete(mexset.exchanges, msgID)
		return true, false
	}

	if _, timedOut := mexset.timeoutExchanges[msgID]; timedOut {
		delete(mexset.timeoutExchanges, msgID)
		return false, true
	}

	return false, false
}

// removeExchange removes a message exchange from the set, if it exists.
// It decrements the sendChRefs wait group, signalling that this exchange no longer has
// any active goroutines that will try to send to sendCh.
func (mexset *messageExchangeSet) removeExchange(msgID uint32) {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("Removing %s message exchange %d", mexset.name, msgID)
	}

	mexset.Lock()
	found, timedOut := mexset.deleteExchange(msgID)
	mexset.Unlock()

	if !found && !timedOut {
		mexset.log.WithFields(
			LogField{"msgID", msgID},
		).Error("Tried to remove exchange multiple times")
		return
	}

	// If the message exchange was found, then we perform clean up actions.
	// These clean up actions can only be run once per exchange.
	mexset.sendChRefs.Done()
	mexset.onRemoved()
}

// timeoutExchange is similar to removeExchange, however it does not decrement
// the sendChRefs wait group, since there could still be a handler running that
// will write to the send channel.
func (mexset *messageExchangeSet) timeoutExchange(msgID uint32) {
	mexset.log.Debugf("Removing %s message exchange %d due to timeout", mexset.name, msgID)

	mexset.Lock()
	delete(mexset.exchanges, msgID)
	mexset.timeoutExchanges[msgID] = struct{}{}
	mexset.Unlock()

	mexset.onRemoved()
}

// waitForSendCh waits for all goroutines with references to sendCh to complete.
func (mexset *messageExchangeSet) waitForSendCh() {
	mexset.sendChRefs.Wait()
}

func (mexset *messageExchangeSet) count() int {
	mexset.RLock()
	count := len(mexset.exchanges)
	mexset.RUnlock()

	return count
}

// forwardPeerFrame forwards a frame from the peer to the appropriate message
// exchange
func (mexset *messageExchangeSet) forwardPeerFrame(frame *Frame) error {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("forwarding %s %s", mexset.name, frame.Header)
	}

	mexset.RLock()
	mex := mexset.exchanges[frame.Header.ID]
	mexset.RUnlock()

	if mex == nil {
		// This is ok since the exchange might have expired or been cancelled
		mexset.log.Infof("received frame %s for %s message exchange that no longer exists",
			frame.Header, mexset.name)
		return nil
	}

	if err := mex.forwardPeerFrame(frame); err != nil {
		mexset.log.Infof("Unable to forward frame %v length %v to %s: %v",
			frame.Header, frame.Header.FrameSize(), mexset.name, err)
		return err
	}

	return nil
}
<commit_msg>Only add to timeoutExchanges if exchange was deleted<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel

import (
	"errors"
	"sync"

	"github.com/uber/tchannel-go/typed"
	"golang.org/x/net/context"
)

var (
	errDuplicateMex        = errors.New("multiple attempts to use the message id")
	errMexChannelFull      = NewSystemError(ErrCodeBusy, "cannot send frame to message exchange channel")
	errUnexpectedFrameType = errors.New("unexpected frame received")
)

const (
	messageExchangeSetInbound  = "inbound"
	messageExchangeSetOutbound = "outbound"

	// mexChannelBufferSize is the size of the message exchange channel buffer.
	mexChannelBufferSize = 2
)

// A messageExchange tracks this Connections's side of a message exchange with a
// peer.  Each message exchange has a channel that can be used to receive
// frames from the peer, and a Context that can controls when the exchange has
// timed out or been cancelled.
type messageExchange struct {
	recvCh    chan *Frame
	ctx       context.Context
	msgID     uint32
	msgType   messageType
	mexset    *messageExchangeSet
	framePool FramePool
}

// forwardPeerFrame forwards a frame from a peer to the message exchange, where
// it can be pulled by whatever application thread is handling the exchange
func (mex *messageExchange) forwardPeerFrame(frame *Frame) error {
	if err := mex.ctx.Err(); err != nil {
		return GetContextError(err)
	}
	select {
	case mex.recvCh <- frame:
		return nil
	case <-mex.ctx.Done():
		// Note: One slow reader processing a large request could stall the connection.
		// If we see this, we need to increase the recvCh buffer size.
		return GetContextError(mex.ctx.Err())
	}
}

// recvPeerFrame waits for a new frame from the peer, or until the context
// expires or is cancelled
func (mex *messageExchange) recvPeerFrame() (*Frame, error) {
	if err := mex.ctx.Err(); err != nil {
		return nil, GetContextError(err)
	}

	select {
	case frame := <-mex.recvCh:
		if frame.Header.ID != mex.msgID {
			mex.mexset.log.WithFields(
				LogField{"msgId", mex.msgID},
				LogField{"header", frame.Header},
			).Error("recvPeerFrame received msg with unexpected ID.")
			return nil, errUnexpectedFrameType
		}
		return frame, nil
	case <-mex.ctx.Done():
		return nil, GetContextError(mex.ctx.Err())
	}
}

// recvPeerFrameOfType waits for a new frame of a given type from the peer, failing
// if the next frame received is not of that type.
// If an error frame is returned, then the errorMessage is returned as the error.
func (mex *messageExchange) recvPeerFrameOfType(msgType messageType) (*Frame, error) {
	frame, err := mex.recvPeerFrame()
	if err != nil {
		return nil, err
	}

	switch frame.Header.messageType {
	case msgType:
		return frame, nil

	case messageTypeError:
		// If we read an error frame, we can release it once we deserialize it.
		defer mex.framePool.Release(frame)

		errMsg := errorMessage{
			id: frame.Header.ID,
		}
		var rbuf typed.ReadBuffer
		rbuf.Wrap(frame.SizedPayload())
		if err := errMsg.read(&rbuf); err != nil {
			return nil, err
		}
		return nil, errMsg

	default:
		// TODO(mmihic): Should be treated as a protocol error
		mex.mexset.log.WithFields(
			LogField{"header", frame.Header},
			LogField{"expectedType", msgType},
			LogField{"expectedID", mex.msgID},
		).Warn("Received unexpected frame.")
		return nil, errUnexpectedFrameType
	}
}

// shutdown shuts down the message exchange, removing it from the message
// exchange set so  that it cannot receive more messages from the peer.  The
// receive channel remains open, however, in case there are concurrent
// goroutines sending to it.
func (mex *messageExchange) shutdown() {
	mex.mexset.removeExchange(mex.msgID)
}

// inboundTimeout is called when an exchange times out, but a handler may still be
// running in the background. Since the handler may still write to the exchange, we
// cannot shutdown the exchange, but we should remove it from the connection's
// exchange list.
func (mex *messageExchange) inboundTimeout() {
	mex.mexset.timeoutExchange(mex.msgID)
}

// A messageExchangeSet manages a set of active message exchanges.  It is
// mainly used to route frames from a peer to the appropriate messageExchange,
// or to cancel or mark a messageExchange as being in error.  Each Connection
// maintains two messageExchangeSets, one to manage exchanges that it has
// initiated (outbound), and another to manage exchanges that the peer has
// initiated (inbound).  The message-type specific handlers are responsible for
// ensuring that their message exchanges are properly registered and removed
// from the corresponding exchange set.
type messageExchangeSet struct {
	sync.RWMutex

	log        Logger
	name       string
	onRemoved  func()
	onAdded    func()
	sendChRefs sync.WaitGroup

	// maps are mutable, and are protected by the mutex.
	exchanges        map[uint32]*messageExchange
	timeoutExchanges map[uint32]struct{}
}

// newMessageExchangeSet creates a new messageExchangeSet with a given name.
func newMessageExchangeSet(log Logger, name string) *messageExchangeSet {
	return &messageExchangeSet{
		name:             name,
		log:              log.WithFields(LogField{"exchange", name}),
		exchanges:        make(map[uint32]*messageExchange),
		timeoutExchanges: make(map[uint32]struct{}),
	}
}

// newExchange creates and adds a new message exchange to this set
func (mexset *messageExchangeSet) newExchange(ctx context.Context, framePool FramePool,
	msgType messageType, msgID uint32, bufferSize int) (*messageExchange, error) {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("Creating new %s message exchange for [%v:%d]", mexset.name, msgType, msgID)
	}

	mex := &messageExchange{
		msgType:   msgType,
		msgID:     msgID,
		ctx:       ctx,
		recvCh:    make(chan *Frame, bufferSize),
		mexset:    mexset,
		framePool: framePool,
	}

	mexset.Lock()
	if existingMex := mexset.exchanges[mex.msgID]; existingMex != nil {
		if existingMex == mex {
			mexset.log.WithFields(
				LogField{"name", mexset.name},
				LogField{"msgType", mex.msgType},
				LogField{"msgID", mex.msgID},
			).Warn("mex registered multiple times.")
		} else {
			mexset.log.WithFields(
				LogField{"msgID", mex.msgID},
				LogField{"existingType", existingMex.msgType},
				LogField{"newType", mex.msgType},
			).Warn("Duplicate msg ID for active and new mex.")
		}

		mexset.Unlock()
		return nil, errDuplicateMex
	}

	mexset.exchanges[mex.msgID] = mex
	mexset.sendChRefs.Add(1)
	mexset.Unlock()

	mexset.onAdded()

	// TODO(mmihic): Put into a deadline ordered heap so we can garbage collected expired exchanges
	return mex, nil
}

// deleteExchange will delete msgID, and return whether it was found or whether it was
// timed out. This method must be called with the lock.
func (mexset *messageExchangeSet) deleteExchange(msgID uint32) (found, timedOut bool) {
	if _, found := mexset.exchanges[msgID]; found {
		delete(mexset.exchanges, msgID)
		return true, false
	}

	if _, timedOut := mexset.timeoutExchanges[msgID]; timedOut {
		delete(mexset.timeoutExchanges, msgID)
		return false, true
	}

	return false, false
}

// removeExchange removes a message exchange from the set, if it exists.
// It decrements the sendChRefs wait group, signalling that this exchange no longer has
// any active goroutines that will try to send to sendCh.
func (mexset *messageExchangeSet) removeExchange(msgID uint32) {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("Removing %s message exchange %d", mexset.name, msgID)
	}

	mexset.Lock()
	found, timedOut := mexset.deleteExchange(msgID)
	mexset.Unlock()

	if !found && !timedOut {
		mexset.log.WithFields(
			LogField{"msgID", msgID},
		).Error("Tried to remove exchange multiple times")
		return
	}

	// If the message exchange was found, then we perform clean up actions.
	// These clean up actions can only be run once per exchange.
	mexset.sendChRefs.Done()
	mexset.onRemoved()
}

// timeoutExchange is similar to removeExchange, however it does not decrement
// the sendChRefs wait group, since there could still be a handler running that
// will write to the send channel.
func (mexset *messageExchangeSet) timeoutExchange(msgID uint32) {
	mexset.log.Debugf("Removing %s message exchange %d due to timeout", mexset.name, msgID)

	mexset.Lock()
	found, timedOut := mexset.deleteExchange(msgID)
	if found || timedOut {
		// Record in timeoutExchange if we deleted the exchange.
		mexset.timeoutExchanges[msgID] = struct{}{}
	}
	mexset.Unlock()

	if timedOut {
		mexset.log.WithFields(LogField{"msgID", msgID}).Error("Exchange timed out already")
	}

	mexset.onRemoved()
}

// waitForSendCh waits for all goroutines with references to sendCh to complete.
func (mexset *messageExchangeSet) waitForSendCh() {
	mexset.sendChRefs.Wait()
}

func (mexset *messageExchangeSet) count() int {
	mexset.RLock()
	count := len(mexset.exchanges)
	mexset.RUnlock()

	return count
}

// forwardPeerFrame forwards a frame from the peer to the appropriate message
// exchange
func (mexset *messageExchangeSet) forwardPeerFrame(frame *Frame) error {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("forwarding %s %s", mexset.name, frame.Header)
	}

	mexset.RLock()
	mex := mexset.exchanges[frame.Header.ID]
	mexset.RUnlock()

	if mex == nil {
		// This is ok since the exchange might have expired or been cancelled
		mexset.log.Infof("received frame %s for %s message exchange that no longer exists",
			frame.Header, mexset.name)
		return nil
	}

	if err := mex.forwardPeerFrame(frame); err != nil {
		mexset.log.Infof("Unable to forward frame %v length %v to %s: %v",
			frame.Header, frame.Header.FrameSize(), mexset.name, err)
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package db

import (
	"database/sql"
	"errors"
	"os"
	"strings"

	// PostgreSQL driver
	_ "github.com/lib/pq"
)

type Maintainer struct {
	GithubHandle string
	Email        string
}

type PortMaintainer struct {
	Primary        *Maintainer
	Others         []*Maintainer
	NoMaintainer   bool
	OpenMaintainer bool
}

type DBHelper interface {
	GetGitHubHandle(email string) (string, error)
	GetPortMaintainer(port string) (*PortMaintainer, error)
}

func NewDBHelper() (DBHelper, error) {
	// TODO: move os.Getenv to main
	tracDB, err := sql.Open("postgres", os.Getenv("TRAC_DB"))
	if err != nil {
		return nil, err
	}
	wwwDB, err := sql.Open("postgres", os.Getenv("WWW_DB"))
	if err != nil {
		return nil, err
	}
	prDB, err := sql.Open("postgres", os.Getenv("PR_DB"))
	if err != nil {
		return nil, err
	}
	return &sqlDBHelper{
		tracDB: tracDB,
		wwwDB:  wwwDB,
		prDB:   prDB,
	}, nil
}

type sqlDBHelper struct {
	tracDB, wwwDB, prDB *sql.DB
}

func (sqlDB *sqlDBHelper) GetGitHubHandle(email string) (string, error) {
	sid := ""
	err := sqlDB.tracDB.QueryRow("SELECT sid "+
		"FROM trac_macports.session_attribute "+
		"WHERE value = $1 "+
		"AND name = 'email' "+
		"AND authenticated = 1 "+
		"LIMIT 1", email).
		Scan(&sid)
	if err != nil {
		return "", err
	}
	return sid, nil
}

// GetPortMaintainer returns the maintainers of a port
func (sqlDB *sqlDBHelper) GetPortMaintainer(port string) (*PortMaintainer, error) {
	rows, err := sqlDB.wwwDB.Query("SELECT maintainer, is_primary "+
		"FROM public.maintainers "+
		"WHERE portfile = $1", port)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	maintainer := new(PortMaintainer)
	maintainerCursor := ""
	isPrimary := false
	rowExist := false

	for rows.Next() {
		if err := rows.Scan(&maintainerCursor, &isPrimary); err != nil {
			return nil, err
		}
		rowExist = true
		switch maintainerCursor {
		case "nomaintainer":
			maintainer.NoMaintainer = true
			continue
		case "openmaintainer":
			maintainer.OpenMaintainer = true
			continue
		}
		if isPrimary {
			maintainer.Primary = sqlDB.parseMaintainer(maintainerCursor)
		} else {
			maintainer.Others = append(maintainer.Others, sqlDB.parseMaintainer(maintainerCursor))
		}
	}

	if !rowExist {
		return nil, errors.New("port not found")
	}

	return maintainer, nil
}

func (sqlDB *sqlDBHelper) parseMaintainer(maintainerFullString string) *Maintainer {
	maintainer := parseMaintainerString(maintainerFullString)
	if maintainer.GithubHandle == "" && maintainer.Email != "" {
		if handle, err := sqlDB.GetGitHubHandle(maintainer.Email); err == nil {
			maintainer.GithubHandle = handle
		}
	}
	return maintainer
}

func parseMaintainerString(maintainerFullString string) *Maintainer {
	maintainerStrings := strings.Split(maintainerFullString, " ")
	maintainer := new(Maintainer)
	for _, maintainerString := range maintainerStrings {
		if strings.HasPrefix(maintainerString, "@") {
			maintainer.GithubHandle = maintainerString[1:]
		} else if strings.Count(maintainerString, ":") == 1 {
			emailParts := strings.Split(maintainerString, ":")
			maintainer.Email = emailParts[1] + "@" + emailParts[0]
		} else {
			maintainer.Email = maintainerString + "@macports.org"
		}
	}
	return maintainer
}
<commit_msg>Establish connection to DB on start<commit_after>package db

import (
	"database/sql"
	"errors"
	"os"
	"strings"

	// PostgreSQL driver
	_ "github.com/lib/pq"
)

type Maintainer struct {
	GithubHandle string
	Email        string
}

type PortMaintainer struct {
	Primary        *Maintainer
	Others         []*Maintainer
	NoMaintainer   bool
	OpenMaintainer bool
}

type DBHelper interface {
	GetGitHubHandle(email string) (string, error)
	GetPortMaintainer(port string) (*PortMaintainer, error)
}

func NewDBHelper() (DBHelper, error) {
	// TODO: move os.Getenv to main
	tracDB, err := sql.Open("postgres", os.Getenv("TRAC_DB"))
	if err != nil {
		return nil, err
	}
	err = tracDB.Ping()
	if err != nil {
		return nil, err
	}

	wwwDB, err := sql.Open("postgres", os.Getenv("WWW_DB"))
	if err != nil {
		return nil, err
	}
	err = wwwDB.Ping()
	if err != nil {
		return nil, err
	}

	prDB, err := sql.Open("postgres", os.Getenv("PR_DB"))
	if err != nil {
		return nil, err
	}
	err = prDB.Ping()
	if err != nil {
		return nil, err
	}

	return &sqlDBHelper{
		tracDB: tracDB,
		wwwDB:  wwwDB,
		prDB:   prDB,
	}, nil
}

type sqlDBHelper struct {
	tracDB, wwwDB, prDB *sql.DB
}

func (sqlDB *sqlDBHelper) GetGitHubHandle(email string) (string, error) {
	sid := ""
	err := sqlDB.tracDB.QueryRow("SELECT sid "+
		"FROM trac_macports.session_attribute "+
		"WHERE value = $1 "+
		"AND name = 'email' "+
		"AND authenticated = 1 "+
		"LIMIT 1", email).
		Scan(&sid)
	if err != nil {
		return "", err
	}
	return sid, nil
}

// GetPortMaintainer returns the maintainers of a port
func (sqlDB *sqlDBHelper) GetPortMaintainer(port string) (*PortMaintainer, error) {
	rows, err := sqlDB.wwwDB.Query("SELECT maintainer, is_primary "+
		"FROM public.maintainers "+
		"WHERE portfile = $1", port)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	maintainer := new(PortMaintainer)
	maintainerCursor := ""
	isPrimary := false
	rowExist := false

	for rows.Next() {
		if err := rows.Scan(&maintainerCursor, &isPrimary); err != nil {
			return nil, err
		}
		rowExist = true
		switch maintainerCursor {
		case "nomaintainer":
			maintainer.NoMaintainer = true
			continue
		case "openmaintainer":
			maintainer.OpenMaintainer = true
			continue
		}
		if isPrimary {
			maintainer.Primary = sqlDB.parseMaintainer(maintainerCursor)
		} else {
			maintainer.Others = append(maintainer.Others, sqlDB.parseMaintainer(maintainerCursor))
		}
	}

	if !rowExist {
		return nil, errors.New("port not found")
	}

	return maintainer, nil
}

func (sqlDB *sqlDBHelper) parseMaintainer(maintainerFullString string) *Maintainer {
	maintainer := parseMaintainerString(maintainerFullString)
	if maintainer.GithubHandle == "" && maintainer.Email != "" {
		if handle, err := sqlDB.GetGitHubHandle(maintainer.Email); err == nil {
			maintainer.GithubHandle = handle
		}
	}
	return maintainer
}

func parseMaintainerString(maintainerFullString string) *Maintainer {
	maintainerStrings := strings.Split(maintainerFullString, " ")
	maintainer := new(Maintainer)
	for _, maintainerString := range maintainerStrings {
		if strings.HasPrefix(maintainerString, "@") {
			maintainer.GithubHandle = maintainerString[1:]
		} else if strings.Count(maintainerString, ":") == 1 {
			emailParts := strings.Split(maintainerString, ":")
			maintainer.Email = emailParts[1] + "@" + emailParts[0]
		} else {
			maintainer.Email = maintainerString + "@macports.org"
		}
	}
	return maintainer
}
<|endoftext|>