<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"github.com/lib/pq"
	"github.com/the42/ogdat"
	"github.com/the42/ogdat/ogdatv21"
	"os"
	"time"
)

type DBer interface {
	Exec(query string, args ...interface{}) (sql.Result, error)
	Prepare(query string) (*sql.Stmt, error)
	Query(query string, args ...interface{}) (*sql.Rows, error)
	QueryRow(query string, args ...interface{}) *sql.Row
}

type DBConn struct {
	DBer
	appid string
}

type State int16

const (
	StateOk State = iota + 1
	StateWarning
	StateError
	StateFatal
)

type DBID int32

func GetDatabaseConnection(appid string) *sql.DB {

	var dburl, dbconnstring string

	if dburl = os.Getenv("DATABASE_URL"); dburl == "" {
		dburl = "postgres://"
	}

	dbconnstring, err := pq.ParseURL(dburl)
	if err != nil {
		fmt.Printf("Invalid Database Url: %s\n", dburl)
		logger.Panicf("Fatal: Invalid Database Url: %s\n", dburl)
	}

	db, err := sql.Open("postgres", dbconnstring)
	if err != nil {
		fmt.Println("Unable to connect to dabase")
		logger.Panicln("Unable to connect to dabase")
	}
	return db
}

func (conn *DBConn) GetLastHit() (*time.Time, error) {
	row := conn.QueryRow("SELECT getlasttimestamp()")

	var t pq.NullTime
	if err := row.Scan(&t); err != nil {
		return nil, err
	}
	if t.Valid {
		return &t.Time, nil
	}
	return nil, nil
}

func (conn DBConn) ResetDatabase() error {
	_, err := conn.Exec("SELECT deleteallentries()")
	if err != nil {
		return err
	}
	return nil
}

func (conn *DBConn) CreateDatabase() error {
	_, err := conn.Exec(postgresdbcreatestatement)
	if err != nil {
		return err
	}
	return nil
}

// Deliberately use no stored procedures
func (conn *DBConn) HeartBeat() error {
	const (
		updatestmt = "UPDATE heartbeat SET ts=$1 WHERE who=$2 AND sysid=$3"
		insertstmt = "INSERT INTO heartbeat(ts, statuscode, statustext, who) VALUES($1, 0, 'Alive', $2)"
	)

	var hbstatement *sql.Stmt
	var sysid DBID

	err := conn.QueryRow("SELECT asi.sysid FROM (SELECT sysid, ts, who, MAX(ts) OVER (PARTITION BY who) max_ts FROM heartbeat) asi WHERE asi.ts = max_ts AND who=$1", conn.appid).Scan(&sysid)

	switch {
	case err == sql.ErrNoRows:
		hbstatement, err = conn.Prepare(insertstmt)
		_, err = hbstatement.Exec(time.Now().UTC(), conn.appid)
	case err != nil:
		return fmt.Errorf("Error heartbeating database: %s", err)
	default:
		hbstatement, err = conn.Prepare(updatestmt)
		_, err = hbstatement.Exec(time.Now().UTC(), conn.appid, sysid)
	}
	defer hbstatement.Close()

	if err != nil {
		return fmt.Errorf("Error executing heartbeat: %s", err)
	}
	return nil
}

// Deliberately use no stored procedures
func (conn DBConn) LogMessage(message string, code State, replacelatest bool) error {

	const (
		updatestmt = "UPDATE heartbeat SET ts=$1, statuscode=$2, statustext=$3 WHERE who=$4 AND sysid=$5"
		insertstmt = "INSERT INTO heartbeat(ts, statuscode, statustext, who) VALUES($1, $2, $3, $4)"
	)

	var hbstatement *sql.Stmt
	var statuscode State
	var sysid DBID

	err := conn.QueryRow("SELECT asi.statuscode, asi.sysid FROM (SELECT sysid, ts, statuscode, who, MAX(ts) OVER (PARTITION BY who) max_ts FROM heartbeat) asi WHERE asi.ts = max_ts AND who=$1", conn.appid).Scan(&statuscode, &sysid)

	switch {
	case err == sql.ErrNoRows:
		hbstatement, err = conn.Prepare(insertstmt)
		_, err = hbstatement.Exec(time.Now().UTC(), code, message, conn.appid)
	case err != nil:
		return fmt.Errorf("Error reading last DBLog status code: %s", err)
	case statuscode != StateOk && replacelatest:
		return fmt.Errorf("Last DBLog caused a non-ok state and update requested, doing nothing")
	case replacelatest:
		hbstatement, err = conn.Prepare(updatestmt)
		_, err = hbstatement.Exec(time.Now().UTC(), code, message, conn.appid, sysid)
	default:
		hbstatement, err = conn.Prepare(insertstmt)
		_, err = hbstatement.Exec(time.Now().UTC(), code, message, conn.appid)
	}
	defer hbstatement.Close()

	if err != nil {
		return fmt.Errorf("Error executing DBLog: %s", err)
	}
	return nil
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func DBStringLen(in string, length int) string {
	return in[:min(length, len(in))]
}

func (conn *DBConn) InsertOrUpdateMetadataInfo(md *ogdatv21.MetaData) (DBID, error) {
	// insertorupdatemetadatainfo(id character varying, pub character varying, cont character varying, descr text, vers character varying, category json, stime timestamp with time zone)
	const stmt = "SELECT * FROM insertorupdatemetadatainfo($1, $2, $3, $4, $5, $6, $7)"

	dbs, err := conn.Prepare(stmt)
	if err != nil {
		return -1, err
	}

	id := DBStringLen(md.Metadata_Identifier.String(), 255)

	pub := md.Publisher
	if pub != nil {
		*pub = DBStringLen(*pub, 255)
	}

	maint := DBStringLen(md.Maintainer_Link.String(), 255)

	desc := md.Description
	if desc != nil {
		*desc = DBStringLen(*desc, 255)
	}

	vers := md.Schema_Name
	if vers != nil {
		*vers = DBStringLen(*vers, 255)
	}

	var cats []string
	if cat := md.Categorization; cat != nil {
		for _, cat := range cat.Kategorie {
			cats = append(cats, cat.ID)
		}
	}
	cat, _ := json.Marshal(cats)

	t := time.Now().UTC()

	row := dbs.QueryRow(id, pub, maint, desc, vers, string(cat), t)

	var sysid DBID
	err = row.Scan(&sysid)
	if err != nil {
		return -1, err
	}
	return sysid, nil
}

func (conn *DBConn) ProtocollCheck(id DBID, messages []ogdat.CheckMessage) error {
	// TODO: decide wheather to insert with a prepare or using a SP
	return nil
}

// Execute Database Timeouting Transaction
func ExecDBTT(conn *sql.DB, timeout time.Duration, statement string, args ...interface{}) (sql.Result, error) {

	tx, err := conn.Begin()
	if err != nil {
		return nil, fmt.Errorf("Unable to create DB Transaction")
	}

	statementreturn := make(chan bool)
	var sqlresult sql.Result
	var execerror error

	go func() {
		sqlresult, execerror = tx.Exec(statement, args)
		statementreturn <- true
	}()

	select {
	case <-statementreturn:
		if execerror != nil {
			tx.Rollback()
		} else {
			execerror = nil
			tx.Commit()
		}

	case <-time.After(timeout):
		tx.Rollback()
		return nil, fmt.Errorf("SQL Statement timed out, rolling back")
	}

	return sqlresult, execerror
}

const postgresdbcreatestatement = `
SET statement_timeout = 0;
SET client_encoding = 'UTF8';
SET standard_conforming_strings = off;
SET check_function_bodies = false;
SET client_min_messages = warning;
SET escape_string_warning = off;
CREATE EXTENSION IF NOT EXISTS plpgsql WITH SCHEMA pg_catalog;
CREATE TYPE odcategory AS ENUM (
    'arbeit',
    'bevölkerung',
    'bildung-und-forschung',
    'finanzen-und-rechnungswesen',
    'geographie-und-planung',
    'gesellschaft-und-soziales',
    'gesundheit',
    'kunst-und-kultur',
    'land-und-forstwirtschaft',
    'sport-und-freizeit',
    'umwelt',
    'verkehr-und-technik',
    'verwaltung-und-politik',
    'wirtschaft-und-tourismus'
);
CREATE TYPE odstatus AS ENUM (
    'updated',
    'inserted',
    'deleted',
    'error_fix',
    'warning_fix',
    'warning',
    'error'
);

CREATE FUNCTION deleteallentries() RETURNS void
    LANGUAGE sql
    AS $$
delete from status;
-- insert into status(reason_text) values('Hallo');
delete from dataset;
$$;

CREATE FUNCTION getlasttimestamp() RETURNS timestamp with time zone
    LANGUAGE sql
    AS $$select max(hittime) from status;$$;

SET default_tablespace = '';

SET default_with_oids = false;

CREATE TABLE dataset (
    sysid integer NOT NULL,
    id character varying(255) NOT NULL,
    publisher character varying(255),
    contact character varying(255) NOT NULL,
    description text,
    version character varying(20) DEFAULT 'v1'::character varying NOT NULL,
    category json
);

CREATE SEQUENCE dataset_sysid_seq
    START WITH 1
    INCREMENT BY 1
    NO MINVALUE
    NO MAXVALUE
    CACHE 1;

ALTER SEQUENCE dataset_sysid_seq OWNED BY dataset.sysid;

SELECT pg_catalog.setval('dataset_sysid_seq', 1, true);

CREATE TABLE heartbeat (
    sysid integer NOT NULL,
    "when" timestamp with time zone,
    statustext character varying(255),
    fetchtime timestamp with time zone,
    statuscode smallint
);

CREATE SEQUENCE heartbeat_sysid_seq
    START WITH 1
    INCREMENT BY 1
    NO MINVALUE
    NO MAXVALUE
    CACHE 1;

ALTER SEQUENCE heartbeat_sysid_seq OWNED BY heartbeat.sysid;

SELECT pg_catalog.setval('heartbeat_sysid_seq', 1, false);

CREATE TABLE status (
    sysid integer NOT NULL,
    datasetid integer NOT NULL,
    reason_text character varying(255),
    field_id integer,
    hittime timestamp with time zone,
    status odstatus
);

CREATE SEQUENCE status_sysid_seq
    START WITH 1
    INCREMENT BY 1
    NO MINVALUE
    NO MAXVALUE
    CACHE 1;

ALTER SEQUENCE status_sysid_seq OWNED BY status.sysid;

SELECT pg_catalog.setval('status_sysid_seq', 1, true);

ALTER TABLE ONLY dataset ALTER COLUMN sysid SET DEFAULT nextval('dataset_sysid_seq'::regclass);

ALTER TABLE ONLY heartbeat ALTER COLUMN sysid SET DEFAULT nextval('heartbeat_sysid_seq'::regclass);

ALTER TABLE ONLY status ALTER COLUMN sysid SET DEFAULT nextval('status_sysid_seq'::regclass);

ALTER TABLE ONLY heartbeat
    ADD CONSTRAINT pk_sysid PRIMARY KEY (sysid);

ALTER TABLE ONLY dataset
    ADD CONSTRAINT pkey PRIMARY KEY (sysid);

ALTER TABLE ONLY status
    ADD CONSTRAINT status_pkey PRIMARY KEY (sysid);

ALTER TABLE ONLY status
    ADD CONSTRAINT status_datasetid_fkey FOREIGN KEY (datasetid) REFERENCES dataset(sysid);
`
<commit_msg>make consistent: receiver is of type pointer type<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"github.com/lib/pq"
	"github.com/the42/ogdat"
	"github.com/the42/ogdat/ogdatv21"
	"os"
	"time"
)

type DBer interface {
	Exec(query string, args ...interface{}) (sql.Result, error)
	Prepare(query string) (*sql.Stmt, error)
	Query(query string, args ...interface{}) (*sql.Rows, error)
	QueryRow(query string, args ...interface{}) *sql.Row
}

type DBConn struct {
	DBer
	appid string
}

type State int16

const (
	StateOk State = iota + 1
	StateWarning
	StateError
	StateFatal
)

type DBID int32

func GetDatabaseConnection(appid string) *sql.DB {

	var dburl, dbconnstring string

	if dburl = os.Getenv("DATABASE_URL"); dburl == "" {
		dburl = "postgres://"
	}

	dbconnstring, err := pq.ParseURL(dburl)
	if err != nil {
		fmt.Printf("Invalid Database Url: %s\n", dburl)
		logger.Panicf("Fatal: Invalid Database Url: %s\n", dburl)
	}

	db, err := sql.Open("postgres", dbconnstring)
	if err != nil {
		fmt.Println("Unable to connect to dabase")
		logger.Panicln("Unable to connect to dabase")
	}
	return db
}

func (conn *DBConn) GetLastHit() (*time.Time, error) {
	row := conn.QueryRow("SELECT getlasttimestamp()")

	var t pq.NullTime
	if err := row.Scan(&t); err != nil {
		return nil, err
	}
	if t.Valid {
		return &t.Time, nil
	}
	return nil, nil
}

func (conn *DBConn) ResetDatabase() error {
	_, err := conn.Exec("SELECT deleteallentries()")
	if err != nil {
		return err
	}
	return nil
}

func (conn *DBConn) CreateDatabase() error {
	_, err := conn.Exec(postgresdbcreatestatement)
	if err != nil {
		return err
	}
	return nil
}

// Deliberately use no stored procedures
func (conn *DBConn) HeartBeat() error {
	const (
		updatestmt = "UPDATE heartbeat SET ts=$1 WHERE who=$2 AND sysid=$3"
		insertstmt = "INSERT INTO heartbeat(ts, statuscode, statustext, who) VALUES($1, 0, 'Alive', $2)"
	)

	var hbstatement *sql.Stmt
	var sysid DBID

	err := conn.QueryRow("SELECT asi.sysid FROM (SELECT sysid, ts, who, MAX(ts) OVER (PARTITION BY who) max_ts FROM heartbeat) asi WHERE asi.ts = max_ts AND who=$1", conn.appid).Scan(&sysid)

	switch {
	case err == sql.ErrNoRows:
		hbstatement, err = conn.Prepare(insertstmt)
		_, err = hbstatement.Exec(time.Now().UTC(), conn.appid)
	case err != nil:
		return fmt.Errorf("Error heartbeating database: %s", err)
	default:
		hbstatement, err = conn.Prepare(updatestmt)
		_, err = hbstatement.Exec(time.Now().UTC(), conn.appid, sysid)
	}
	defer hbstatement.Close()

	if err != nil {
		return fmt.Errorf("Error executing heartbeat: %s", err)
	}
	return nil
}

// Deliberately use no stored procedures
func (conn *DBConn) LogMessage(message string, code State, replacelatest bool) error {

	const (
		updatestmt = "UPDATE heartbeat SET ts=$1, statuscode=$2, statustext=$3 WHERE who=$4 AND sysid=$5"
		insertstmt = "INSERT INTO heartbeat(ts, statuscode, statustext, who) VALUES($1, $2, $3, $4)"
	)

	var hbstatement *sql.Stmt
	var statuscode State
	var sysid DBID

	err := conn.QueryRow("SELECT asi.statuscode, asi.sysid FROM (SELECT sysid, ts, statuscode, who, MAX(ts) OVER (PARTITION BY who) max_ts FROM heartbeat) asi WHERE asi.ts = max_ts AND who=$1", conn.appid).Scan(&statuscode, &sysid)

	switch {
	case err == sql.ErrNoRows:
		hbstatement, err = conn.Prepare(insertstmt)
		_, err = hbstatement.Exec(time.Now().UTC(), code, message, conn.appid)
	case err != nil:
		return fmt.Errorf("Error reading last DBLog status code: %s", err)
	case statuscode != StateOk && replacelatest:
		return fmt.Errorf("Last DBLog caused a non-ok state and update requested, doing nothing")
	case replacelatest:
		hbstatement, err = conn.Prepare(updatestmt)
		_, err = hbstatement.Exec(time.Now().UTC(), code, message, conn.appid, sysid)
	default:
		hbstatement, err = conn.Prepare(insertstmt)
		_, err = hbstatement.Exec(time.Now().UTC(), code, message, conn.appid)
	}
	defer hbstatement.Close()

	if err != nil {
		return fmt.Errorf("Error executing DBLog: %s", err)
	}
	return nil
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func DBStringLen(in string, length int) string {
	return in[:min(length, len(in))]
}

func (conn *DBConn) InsertOrUpdateMetadataInfo(md *ogdatv21.MetaData) (DBID, error) {
	// insertorupdatemetadatainfo(id character varying, pub character varying, cont character varying, descr text, vers character varying, category json, stime timestamp with time zone)
	const stmt = "SELECT * FROM insertorupdatemetadatainfo($1, $2, $3, $4, $5, $6, $7)"

	dbs, err := conn.Prepare(stmt)
	if err != nil {
		return -1, err
	}

	id := DBStringLen(md.Metadata_Identifier.String(), 255)

	pub := md.Publisher
	if pub != nil {
		*pub = DBStringLen(*pub, 255)
	}

	maint := DBStringLen(md.Maintainer_Link.String(), 255)

	desc := md.Description
	if desc != nil {
		*desc = DBStringLen(*desc, 255)
	}

	vers := md.Schema_Name
	if vers != nil {
		*vers = DBStringLen(*vers, 255)
	}

	var cats []string
	if cat := md.Categorization; cat != nil {
		for _, cat := range cat.Kategorie {
			cats = append(cats, cat.ID)
		}
	}
	cat, _ := json.Marshal(cats)

	t := time.Now().UTC()

	row := dbs.QueryRow(id, pub, maint, desc, vers, string(cat), t)

	var sysid DBID
	err = row.Scan(&sysid)
	if err != nil {
		return -1, err
	}
	return sysid, nil
}

func (conn *DBConn) ProtocollCheck(id DBID, messages []ogdat.CheckMessage) error {
	// TODO: decide wheather to insert with a prepare or using a SP
	return nil
}

// Execute Database Timeouting Transaction
func ExecDBTT(conn *sql.DB, timeout time.Duration, statement string, args ...interface{}) (sql.Result, error) {

	tx, err := conn.Begin()
	if err != nil {
		return nil, fmt.Errorf("Unable to create DB Transaction")
	}

	statementreturn := make(chan bool)
	var sqlresult sql.Result
	var execerror error

	go func() {
		sqlresult, execerror = tx.Exec(statement, args)
		statementreturn <- true
	}()

	select {
	case <-statementreturn:
		if execerror != nil {
			tx.Rollback()
		} else {
			execerror = nil
			tx.Commit()
		}

	case <-time.After(timeout):
		tx.Rollback()
		return nil, fmt.Errorf("SQL Statement timed out, rolling back")
	}

	return sqlresult, execerror
}

const postgresdbcreatestatement = `
SET statement_timeout = 0;
SET client_encoding = 'UTF8';
SET standard_conforming_strings = off;
SET check_function_bodies = false;
SET client_min_messages = warning;
SET escape_string_warning = off;
CREATE EXTENSION IF NOT EXISTS plpgsql WITH SCHEMA pg_catalog;
CREATE TYPE odcategory AS ENUM (
    'arbeit',
    'bevölkerung',
    'bildung-und-forschung',
    'finanzen-und-rechnungswesen',
    'geographie-und-planung',
    'gesellschaft-und-soziales',
    'gesundheit',
    'kunst-und-kultur',
    'land-und-forstwirtschaft',
    'sport-und-freizeit',
    'umwelt',
    'verkehr-und-technik',
    'verwaltung-und-politik',
    'wirtschaft-und-tourismus'
);
CREATE TYPE odstatus AS ENUM (
    'updated',
    'inserted',
    'deleted',
    'error_fix',
    'warning_fix',
    'warning',
    'error'
);

CREATE FUNCTION deleteallentries() RETURNS void
    LANGUAGE sql
    AS $$
delete from status;
-- insert into status(reason_text) values('Hallo');
delete from dataset;
$$;

CREATE FUNCTION getlasttimestamp() RETURNS timestamp with time zone
    LANGUAGE sql
    AS $$select max(hittime) from status;$$;

SET default_tablespace = '';

SET default_with_oids = false;

CREATE TABLE dataset (
    sysid integer NOT NULL,
    id character varying(255) NOT NULL,
    publisher character varying(255),
    contact character varying(255) NOT NULL,
    description text,
    version character varying(20) DEFAULT 'v1'::character varying NOT NULL,
    category json
);

CREATE SEQUENCE dataset_sysid_seq
    START WITH 1
    INCREMENT BY 1
    NO MINVALUE
    NO MAXVALUE
    CACHE 1;

ALTER SEQUENCE dataset_sysid_seq OWNED BY dataset.sysid;

SELECT pg_catalog.setval('dataset_sysid_seq', 1, true);

CREATE TABLE heartbeat (
    sysid integer NOT NULL,
    "when" timestamp with time zone,
    statustext character varying(255),
    fetchtime timestamp with time zone,
    statuscode smallint
);

CREATE SEQUENCE heartbeat_sysid_seq
    START WITH 1
    INCREMENT BY 1
    NO MINVALUE
    NO MAXVALUE
    CACHE 1;

ALTER SEQUENCE heartbeat_sysid_seq OWNED BY heartbeat.sysid;

SELECT pg_catalog.setval('heartbeat_sysid_seq', 1, false);

CREATE TABLE status (
    sysid integer NOT NULL,
    datasetid integer NOT NULL,
    reason_text character varying(255),
    field_id integer,
    hittime timestamp with time zone,
    status odstatus
);

CREATE SEQUENCE status_sysid_seq
    START WITH 1
    INCREMENT BY 1
    NO MINVALUE
    NO MAXVALUE
    CACHE 1;

ALTER SEQUENCE status_sysid_seq OWNED BY status.sysid;

SELECT pg_catalog.setval('status_sysid_seq', 1, true);

ALTER TABLE ONLY dataset ALTER COLUMN sysid SET DEFAULT nextval('dataset_sysid_seq'::regclass);

ALTER TABLE ONLY heartbeat ALTER COLUMN sysid SET DEFAULT nextval('heartbeat_sysid_seq'::regclass);

ALTER TABLE ONLY status ALTER COLUMN sysid SET DEFAULT nextval('status_sysid_seq'::regclass);

ALTER TABLE ONLY heartbeat
    ADD CONSTRAINT pk_sysid PRIMARY KEY (sysid);

ALTER TABLE ONLY dataset
    ADD CONSTRAINT pkey PRIMARY KEY (sysid);

ALTER TABLE ONLY status
    ADD CONSTRAINT status_pkey PRIMARY KEY (sysid);

ALTER TABLE ONLY status
    ADD CONSTRAINT status_datasetid_fkey FOREIGN KEY (datasetid) REFERENCES dataset(sysid);
`
<|endoftext|>


<commit_before>package command

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"os/user"
	"strings"

	"github.com/hashicorp/vault/builtin/logical/ssh"
	"github.com/hashicorp/vault/meta"
	"github.com/mitchellh/mapstructure"
)

// SSHCommand is a Command that establishes a SSH connection
// with target by generating a dynamic key
type SSHCommand struct {
	meta.Meta
}

// Structure to hold the fields returned when asked for a credential from SSHh backend.
type SSHCredentialResp struct {
	KeyType  string `mapstructure:"key_type"`
	Key      string `mapstructure:"key"`
	Username string `mapstructure:"username"`
	IP       string `mapstructure:"ip"`
	Port     string `mapstructure:"port"`
}

func (c *SSHCommand) Run(args []string) int {
	var role, mountPoint, format, userKnownHostsFile, strictHostKeyChecking string
	var noExec bool
	var sshCmdArgs []string
	var sshDynamicKeyFileName string
	flags := c.Meta.FlagSet("ssh", meta.FlagSetDefault)
	flags.StringVar(&strictHostKeyChecking, "strict-host-key-checking", "", "")
	flags.StringVar(&userKnownHostsFile, "user-known-hosts-file", "", "")
	flags.StringVar(&format, "format", "table", "")
	flags.StringVar(&role, "role", "", "")
	flags.StringVar(&mountPoint, "mount-point", "ssh", "")
	flags.BoolVar(&noExec, "no-exec", false, "")

	flags.Usage = func() { c.Ui.Error(c.Help()) }
	if err := flags.Parse(args); err != nil {
		return 1
	}

	// If the flag is already set then it takes the precedence. If the flag is not
	// set, try setting it from env var.
	if os.Getenv("VAULT_SSH_STRICT_HOST_KEY_CHECKING") != "" && strictHostKeyChecking == "" {
		strictHostKeyChecking = os.Getenv("VAULT_SSH_STRICT_HOST_KEY_CHECKING")
	}
	// Assign default value if both flag and env var are not set
	if strictHostKeyChecking == "" {
		strictHostKeyChecking = "ask"
	}

	// If the flag is already set then it takes the precedence. If the flag is not
	// set, try setting it from env var.
	if os.Getenv("VAULT_SSH_USER_KNOWN_HOSTS_FILE") != "" && userKnownHostsFile == "" {
		userKnownHostsFile = os.Getenv("VAULT_SSH_USER_KNOWN_HOSTS_FILE")
	}
	// Assign default value if both flag and env var are not set
	if userKnownHostsFile == "" {
		userKnownHostsFile = "~/.ssh/known_hosts"
	}

	args = flags.Args()
	if len(args) < 1 {
		c.Ui.Error("ssh expects at least one argument")
		return 1
	}

	client, err := c.Client()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
		return 1
	}

	// split the parameter username@ip
	input := strings.Split(args[0], "@")
	var username string
	var ipAddr string

	// If only IP is mentioned and username is skipped, assume username to
	// be the current username. Vault SSH role's default username could have
	// been used, but in order to retain the consistency with SSH command,
	// current username is employed.
	if len(input) == 1 {
		u, err := user.Current()
		if err != nil {
			c.Ui.Error(fmt.Sprintf("Error fetching username: %s", err))
			return 1
		}
		username = u.Username
		ipAddr = input[0]
	} else if len(input) == 2 {
		username = input[0]
		ipAddr = input[1]
	} else {
		c.Ui.Error(fmt.Sprintf("Invalid parameter: %s", args[0]))
		return 1
	}

	// Resolving domain names to IP address on the client side.
	// Vault only deals with IP addresses.
	ip, err := net.ResolveIPAddr("ip", ipAddr)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error resolving IP Address: %s", err))
		return 1
	}

	// Credentials are generated only against a registered role. If user
	// does not specify a role with the SSH command, then lookup API is used
	// to fetch all the roles with which this IP is associated. If there is
	// only one role associated with it, use it to establish the connection.
	if role == "" {
		role, err = c.defaultRole(mountPoint, ip.String())
		if err != nil {
			c.Ui.Error(fmt.Sprintf("Error choosing role: %s", err))
			return 1
		}
		// Print the default role chosen so that user knows the role name
		// if something doesn't work. If the role chosen is not allowed to
		// be used by the user (ACL enforcement), then user should see an
		// error message accordingly.
		c.Ui.Output(fmt.Sprintf("Vault SSH: Role: %s", role))
	}

	data := map[string]interface{}{
		"username": username,
		"ip":       ip.String(),
	}

	keySecret, err := client.SSHWithMountPoint(mountPoint).Credential(role, data)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error getting key for SSH session:%s", err))
		return 1
	}

	// if no-exec was chosen, just print out the secret and return.
	if noExec {
		return OutputSecret(c.Ui, format, keySecret)
	}

	// Port comes back as a json.Number which mapstructure doesn't like, so convert it
	if keySecret.Data["port"] != nil {
		keySecret.Data["port"] = keySecret.Data["port"].(json.Number).String()
	}
	var resp SSHCredentialResp
	if err := mapstructure.Decode(keySecret.Data, &resp); err != nil {
		c.Ui.Error(fmt.Sprintf("Error parsing the credential response:%s", err))
		return 1
	}

	if resp.KeyType == ssh.KeyTypeDynamic {
		if len(resp.Key) == 0 {
			c.Ui.Error(fmt.Sprintf("Invalid key"))
			return 1
		}
		sshDynamicKeyFileName = fmt.Sprintf("vault_ssh_%s_%s", username, ip.String())
		err = ioutil.WriteFile(sshDynamicKeyFileName, []byte(resp.Key), 0600)
		sshCmdArgs = append(sshCmdArgs, []string{"-i", sshDynamicKeyFileName}...)

	} else if resp.KeyType == ssh.KeyTypeOTP {
		// Check if the application 'sshpass' is installed in the client machine.
		// If it is then, use it to automate typing in OTP to the prompt. Unfortunately,
		// it was not possible to automate it without a third-party application, with
		// only the Go libraries.
		// Feel free to try and remove this dependency.
		sshpassPath, err := exec.LookPath("sshpass")
		if err == nil {
			sshCmdArgs = append(sshCmdArgs, []string{"-p", string(resp.Key), "ssh", "-o UserKnownHostsFile=" + userKnownHostsFile, "-o StrictHostKeyChecking=" + strictHostKeyChecking, "-p", resp.Port, username + "@" + ip.String()}...)
			sshCmd := exec.Command(sshpassPath, sshCmdArgs...)
			sshCmd.Stdin = os.Stdin
			sshCmd.Stdout = os.Stdout
			err = sshCmd.Run()
			if err != nil {
				c.Ui.Error(fmt.Sprintf("Failed to establish SSH connection:%s", err))
			}
			return 0
		}
		c.Ui.Output("OTP for the session is " + resp.Key)
		c.Ui.Output("[Note: Install 'sshpass' to automate typing in OTP]")
	}
	sshCmdArgs = append(sshCmdArgs, []string{"-o UserKnownHostsFile=" + userKnownHostsFile, "-o StrictHostKeyChecking=" + strictHostKeyChecking, "-p", resp.Port, username + "@" + ip.String()}...)

	sshCmd := exec.Command("ssh", sshCmdArgs...)
	sshCmd.Stdin = os.Stdin
	sshCmd.Stdout = os.Stdout

	// Running the command as a separate command. The reason for using exec.Command instead
	// of using crypto/ssh package is that, this way, user can have the same feeling of
	// connecting to remote hosts with the ssh command. Package crypto/ssh did not have a way
	// to establish an independent session like this.
	err = sshCmd.Run()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error while running ssh command:%s", err))
	}

	// Delete the temporary key file generated by the command.
	if resp.KeyType == ssh.KeyTypeDynamic {
		// Ignoring the error from the below call since it is not a security
		// issue if the deletion of file is not successful. User is authorized
		// to have this secret.
		os.Remove(sshDynamicKeyFileName)
	}

	// If the session established was longer than the lease expiry, the secret
	// might have been revoked already. If not, then revoke it. Since the key
	// file is deleted and since user doesn't know the credential anymore, there
	// is not point in Vault maintaining this secret anymore. Every time the command
	// is run, a fresh credential is generated anyways.
	err = client.Sys().Revoke(keySecret.LeaseID)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error revoking the key: %s", err))
	}

	return 0
}

// If user did not provide the role with which SSH connection has
// to be established and if there is only one role associated with
// the IP, it is used by default.
func (c *SSHCommand) defaultRole(mountPoint, ip string) (string, error) {
	data := map[string]interface{}{
		"ip": ip,
	}
	client, err := c.Client()
	if err != nil {
		return "", err
	}
	secret, err := client.Logical().Write(mountPoint+"/lookup", data)
	if err != nil {
		return "", fmt.Errorf("Error finding roles for IP %s: %s", ip, err)

	}
	if secret == nil {
		return "", fmt.Errorf("Error finding roles for IP %s: %s", ip, err)
	}

	if secret.Data["roles"] == nil {
		return "", fmt.Errorf("No matching roles found for IP %s", ip)
	}

	if len(secret.Data["roles"].([]interface{})) == 1 {
		return secret.Data["roles"].([]interface{})[0].(string), nil
	} else {
		var roleNames string
		for _, item := range secret.Data["roles"].([]interface{}) {
			roleNames += item.(string) + ", "
		}
		roleNames = strings.TrimRight(roleNames, ", ")
		return "", fmt.Errorf("Roles:[%s]"+`
		Multiple roles are registered for this IP.
		Select a role using '-role' option.
		Note that all roles may not be permitted, based on ACLs.`, roleNames)
	}
}

func (c *SSHCommand) Synopsis() string {
	return "Initiate a SSH session"
}

func (c *SSHCommand) Help() string {
	helpText := `
Usage: vault ssh [options] username@ip

  Establishes an SSH connection with the target machine.

  This command generates a key and uses it to establish an SSH
  connection with the target machine. This operation requires
  that SSH backend is mounted and at least one 'role' be registed
  with vault at priori.

  For setting up SSH backends with one-time-passwords, installation
  of agent in target machines is required.
  See [https://github.com/hashicorp/vault-ssh-agent]

General Options:
` + meta.GeneralOptionsUsage() + `
SSH Options:

	-role				Role to be used to create the key.
					Each IP is associated with a role. To see the associated
					roles with IP, use "lookup" endpoint. If you are certain
					that there is only one role associated with the IP, you can
					skip mentioning the role. It will be chosen by default.  If
					there are no roles associated with the IP, register the
					CIDR block of that IP using the "roles/" endpoint.

	-no-exec			Shows the credentials but does not establish connection.

	-mount-point			Mount point of SSH backend. If the backend is mounted at
					'ssh', which is the default as well, this parameter can be
					skipped.

	-format				If no-exec option is enabled, then the credentials will be
					printed out and SSH connection will not be established. The
					format of the output can be 'json' or 'table'. JSON output
					is useful when writing scripts. Default is 'table'.

	-strict-host-key-checking	This option corresponds to StrictHostKeyChecking of SSH configuration.
					If 'sshpass' is employed to enable automated login, then if host key
					is not "known" to the client, 'vault ssh' command will fail. Set this
					option to "no" to bypass the host key checking. Defaults to "ask".
					Can also be specified with VAULT_SSH_STRICT_HOST_KEY_CHECKING environment
					variable.

	-user-known-hosts-file		This option corresponds to UserKnownHostsFile of SSH configuration.
					Assigns the file to use for storing the host keys. If this option is
					set to "/dev/null" along with "-strict-host-key-checking=no", both
					warnings and host key checking can be avoided while establishing the
					connection. Defaults to "~/.ssh/known_hosts". Can also be specified
					with VAULT_SSH_USER_KNOWN_HOSTS_FILE environment variable.
`
	return strings.TrimSpace(helpText)
}
<commit_msg>Support execution of remote commands using 'vault ssh'<commit_after>package command

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"os/user"
	"strings"

	"github.com/hashicorp/vault/builtin/logical/ssh"
	"github.com/hashicorp/vault/meta"
	"github.com/mitchellh/mapstructure"
)

// SSHCommand is a Command that establishes a SSH connection
// with target by generating a dynamic key
type SSHCommand struct {
	meta.Meta
}

// Structure to hold the fields returned when asked for a credential from SSHh backend.
type SSHCredentialResp struct {
	KeyType  string `mapstructure:"key_type"`
	Key      string `mapstructure:"key"`
	Username string `mapstructure:"username"`
	IP       string `mapstructure:"ip"`
	Port     string `mapstructure:"port"`
}

func (c *SSHCommand) Run(args []string) int {
	var role, mountPoint, format, userKnownHostsFile, strictHostKeyChecking string
	var noExec bool
	var sshCmdArgs []string
	var sshDynamicKeyFileName string
	flags := c.Meta.FlagSet("ssh", meta.FlagSetDefault)
	flags.StringVar(&strictHostKeyChecking, "strict-host-key-checking", "", "")
	flags.StringVar(&userKnownHostsFile, "user-known-hosts-file", "", "")
	flags.StringVar(&format, "format", "table", "")
	flags.StringVar(&role, "role", "", "")
	flags.StringVar(&mountPoint, "mount-point", "ssh", "")
	flags.BoolVar(&noExec, "no-exec", false, "")

	flags.Usage = func() { c.Ui.Error(c.Help()) }
	if err := flags.Parse(args); err != nil {
		return 1
	}

	// If the flag is already set then it takes the precedence. If the flag is not
	// set, try setting it from env var.
	if os.Getenv("VAULT_SSH_STRICT_HOST_KEY_CHECKING") != "" && strictHostKeyChecking == "" {
		strictHostKeyChecking = os.Getenv("VAULT_SSH_STRICT_HOST_KEY_CHECKING")
	}
	// Assign default value if both flag and env var are not set
	if strictHostKeyChecking == "" {
		strictHostKeyChecking = "ask"
	}

	// If the flag is already set then it takes the precedence. If the flag is not
	// set, try setting it from env var.
	if os.Getenv("VAULT_SSH_USER_KNOWN_HOSTS_FILE") != "" && userKnownHostsFile == "" {
		userKnownHostsFile = os.Getenv("VAULT_SSH_USER_KNOWN_HOSTS_FILE")
	}
	// Assign default value if both flag and env var are not set
	if userKnownHostsFile == "" {
		userKnownHostsFile = "~/.ssh/known_hosts"
	}

	args = flags.Args()
	if len(args) < 1 {
		c.Ui.Error("ssh expects at least one argument")
		return 1
	}

	client, err := c.Client()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
		return 1
	}

	// split the parameter username@ip
	input := strings.Split(args[0], "@")
	var username string
	var ipAddr string

	// If only IP is mentioned and username is skipped, assume username to
	// be the current username. Vault SSH role's default username could have
	// been used, but in order to retain the consistency with SSH command,
	// current username is employed.
	if len(input) == 1 {
		u, err := user.Current()
		if err != nil {
			c.Ui.Error(fmt.Sprintf("Error fetching username: %s", err))
			return 1
		}
		username = u.Username
		ipAddr = input[0]
	} else if len(input) == 2 {
		username = input[0]
		ipAddr = input[1]
	} else {
		c.Ui.Error(fmt.Sprintf("Invalid parameter: %s", args[0]))
		return 1
	}

	// Resolving domain names to IP address on the client side.
	// Vault only deals with IP addresses.
	ip, err := net.ResolveIPAddr("ip", ipAddr)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error resolving IP Address: %s", err))
		return 1
	}

	// Credentials are generated only against a registered role. If user
	// does not specify a role with the SSH command, then lookup API is used
	// to fetch all the roles with which this IP is associated. If there is
	// only one role associated with it, use it to establish the connection.
	if role == "" {
		role, err = c.defaultRole(mountPoint, ip.String())
		if err != nil {
			c.Ui.Error(fmt.Sprintf("Error choosing role: %s", err))
			return 1
		}
		// Print the default role chosen so that user knows the role name
		// if something doesn't work. If the role chosen is not allowed to
		// be used by the user (ACL enforcement), then user should see an
		// error message accordingly.
		c.Ui.Output(fmt.Sprintf("Vault SSH: Role: %s", role))
	}

	data := map[string]interface{}{
		"username": username,
		"ip":       ip.String(),
	}

	keySecret, err := client.SSHWithMountPoint(mountPoint).Credential(role, data)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error getting key for SSH session:%s", err))
		return 1
	}

	// if no-exec was chosen, just print out the secret and return.
	if noExec {
		return OutputSecret(c.Ui, format, keySecret)
	}

	// Port comes back as a json.Number which mapstructure doesn't like, so convert it
	if keySecret.Data["port"] != nil {
		keySecret.Data["port"] = keySecret.Data["port"].(json.Number).String()
	}
	var resp SSHCredentialResp
	if err := mapstructure.Decode(keySecret.Data, &resp); err != nil {
		c.Ui.Error(fmt.Sprintf("Error parsing the credential response:%s", err))
		return 1
	}

	if resp.KeyType == ssh.KeyTypeDynamic {
		if len(resp.Key) == 0 {
			c.Ui.Error(fmt.Sprintf("Invalid key"))
			return 1
		}
		sshDynamicKeyFileName = fmt.Sprintf("vault_ssh_%s_%s", username, ip.String())
		err = ioutil.WriteFile(sshDynamicKeyFileName, []byte(resp.Key), 0600)
		sshCmdArgs = append(sshCmdArgs, []string{"-i", sshDynamicKeyFileName}...)

	} else if resp.KeyType == ssh.KeyTypeOTP {
		// Check if the application 'sshpass' is installed in the client machine.
		// If it is then, use it to automate typing in OTP to the prompt. Unfortunately,
		// it was not possible to automate it without a third-party application, with
		// only the Go libraries.
		// Feel free to try and remove this dependency.
		sshpassPath, err := exec.LookPath("sshpass")
		if err == nil {
			sshCmdArgs = append(sshCmdArgs, []string{"-p", string(resp.Key), "ssh", "-o UserKnownHostsFile=" + userKnownHostsFile, "-o StrictHostKeyChecking=" + strictHostKeyChecking, "-p", resp.Port, username + "@" + ip.String()}...)
			if len(args) > 1 {
				sshCmdArgs = append(sshCmdArgs, args[1:]...)
			}
			sshCmd := exec.Command(sshpassPath, sshCmdArgs...)
			sshCmd.Stdin = os.Stdin
			sshCmd.Stdout = os.Stdout
			err = sshCmd.Run()
			if err != nil {
				c.Ui.Error(fmt.Sprintf("Failed to establish SSH connection:%s", err))
			}
			return 0
		}
		c.Ui.Output("OTP for the session is " + resp.Key)
		c.Ui.Output("[Note: Install 'sshpass' to automate typing in OTP]")
	}
	sshCmdArgs = append(sshCmdArgs, []string{"-o UserKnownHostsFile=" + userKnownHostsFile, "-o StrictHostKeyChecking=" + strictHostKeyChecking, "-p", resp.Port, username + "@" + ip.String()}...)
	if len(args) > 1 {
		sshCmdArgs = append(sshCmdArgs, args[1:]...)
	}

	sshCmd := exec.Command("ssh", sshCmdArgs...)
	sshCmd.Stdin = os.Stdin
	sshCmd.Stdout = os.Stdout

	// Running the command as a separate command. The reason for using exec.Command instead
	// of using crypto/ssh package is that, this way, user can have the same feeling of
	// connecting to remote hosts with the ssh command. Package crypto/ssh did not have a way
	// to establish an independent session like this.
	err = sshCmd.Run()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error while running ssh command:%s", err))
	}

	// Delete the temporary key file generated by the command.
	if resp.KeyType == ssh.KeyTypeDynamic {
		// Ignoring the error from the below call since it is not a security
		// issue if the deletion of file is not successful. User is authorized
		// to have this secret.
		os.Remove(sshDynamicKeyFileName)
	}

	// If the session established was longer than the lease expiry, the secret
	// might have been revoked already. If not, then revoke it. Since the key
	// file is deleted and since user doesn't know the credential anymore, there
	// is not point in Vault maintaining this secret anymore. Every time the command
	// is run, a fresh credential is generated anyways.
	err = client.Sys().Revoke(keySecret.LeaseID)
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error revoking the key: %s", err))
	}

	return 0
}

// If user did not provide the role with which SSH connection has
// to be established and if there is only one role associated with
// the IP, it is used by default.
func (c *SSHCommand) defaultRole(mountPoint, ip string) (string, error) {
	data := map[string]interface{}{
		"ip": ip,
	}
	client, err := c.Client()
	if err != nil {
		return "", err
	}
	secret, err := client.Logical().Write(mountPoint+"/lookup", data)
	if err != nil {
		return "", fmt.Errorf("Error finding roles for IP %s: %s", ip, err)

	}
	if secret == nil {
		return "", fmt.Errorf("Error finding roles for IP %s: %s", ip, err)
	}

	if secret.Data["roles"] == nil {
		return "", fmt.Errorf("No matching roles found for IP %s", ip)
	}

	if len(secret.Data["roles"].([]interface{})) == 1 {
		return secret.Data["roles"].([]interface{})[0].(string), nil
	} else {
		var roleNames string
		for _, item := range secret.Data["roles"].([]interface{}) {
			roleNames += item.(string) + ", "
		}
		roleNames = strings.TrimRight(roleNames, ", ")
		return "", fmt.Errorf("Roles:[%s]"+`
		Multiple roles are registered for this IP.
		Select a role using '-role' option.
		Note that all roles may not be permitted, based on ACLs.`, roleNames)
	}
}

func (c *SSHCommand) Synopsis() string {
	return "Initiate a SSH session"
}

func (c *SSHCommand) Help() string {
	helpText := `
Usage: vault ssh [options] username@ip

  Establishes an SSH connection with the target machine.

  This command generates a key and uses it to establish an SSH
  connection with the target machine. This operation requires
  that SSH backend is mounted and at least one 'role' be registed
  with vault at priori.

  For setting up SSH backends with one-time-passwords, installation
  of agent in target machines is required.
  See [https://github.com/hashicorp/vault-ssh-agent]

General Options:
` + meta.GeneralOptionsUsage() + `
SSH Options:

	-role				Role to be used to create the key.
					Each IP is associated with a role. To see the associated
					roles with IP, use "lookup" endpoint. If you are certain
					that there is only one role associated with the IP, you can
					skip mentioning the role. It will be chosen by default.  If
					there are no roles associated with the IP, register the
					CIDR block of that IP using the "roles/" endpoint.

	-no-exec			Shows the credentials but does not establish connection.

	-mount-point			Mount point of SSH backend. If the backend is mounted at
					'ssh', which is the default as well, this parameter can be
					skipped.

	-format				If no-exec option is enabled, then the credentials will be
					printed out and SSH connection will not be established. The
					format of the output can be 'json' or 'table'. JSON output
					is useful when writing scripts. Default is 'table'.

	-strict-host-key-checking	This option corresponds to StrictHostKeyChecking of SSH configuration.
					If 'sshpass' is employed to enable automated login, then if host key
					is not "known" to the client, 'vault ssh' command will fail. Set this
					option to "no" to bypass the host key checking. Defaults to "ask".
					Can also be specified with VAULT_SSH_STRICT_HOST_KEY_CHECKING environment
					variable.

	-user-known-hosts-file		This option corresponds to UserKnownHostsFile of SSH configuration.
					Assigns the file to use for storing the host keys. If this option is
					set to "/dev/null" along with "-strict-host-key-checking=no", both
					warnings and host key checking can be avoided while establishing the
					connection. Defaults to "~/.ssh/known_hosts". Can also be specified
					with VAULT_SSH_USER_KNOWN_HOSTS_FILE environment variable.
`
	return strings.TrimSpace(helpText)
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Workiva, LLC

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

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

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

package batcher

import (
	"errors"
	"sync"
	"sync/atomic"
	"time"
)

// Batcher provides an API for accumulating items into a batch for processing.
type Batcher interface {
	// Put adds items to the batcher.
	Put(interface{}) error

	// Get retrieves a batch from the batcher. This call will block until
	// one of the conditions for a "complete" batch is reached.
	Get() ([]interface{}, error)

	// Flush forcibly completes the batch currently being built
	Flush() error

	// Dispose will dispose of the batcher. Any calls to Put or Flush
	// will return ErrDisposed, calls to Get will return an error iff
	// there are no more ready batches.
	Dispose()

	// IsDisposed will determine if the batcher is disposed
	IsDisposed() bool
}

// ErrDisposed is the error returned for a disposed Batcher
var ErrDisposed = errors.New("batcher: disposed")

// CalculateBytes evaluates the number of bytes in an item added to a Batcher.
type CalculateBytes func(interface{}) uint

type basicBatcher struct {
	maxTime        time.Duration
	maxItems       uint
	maxBytes       uint
	calculateBytes CalculateBytes
	disposed       bool
	items          []interface{}
	lock           sync.RWMutex
	batchChan      chan []interface{}
	availableBytes uint
	waiting        int32
}

// New creates a new Batcher using the provided arguments.
// Batch readiness can be determined in three ways:
//   - Maximum number of bytes per batch
//   - Maximum number of items per batch
//   - Maximum amount of time waiting for a batch
// Values of zero for one of these fields indicate they should not be
// taken into account when evaluating the readiness of a batch.
func New(maxTime time.Duration, maxItems, maxBytes, queueLen uint, calculate CalculateBytes) (Batcher, error) {
	if maxBytes > 0 && calculate == nil {
		return nil, errors.New("batcher: must provide CalculateBytes function")
	}

	return &basicBatcher{
		maxTime:        maxTime,
		maxItems:       maxItems,
		maxBytes:       maxBytes,
		calculateBytes: calculate,
		items:          make([]interface{}, 0, maxItems),
		batchChan:      make(chan []interface{}, queueLen),
	}, nil
}

// Put adds items to the batcher. If Put is continually called without calls to
// Get, an unbounded number of go-routines will be generated.
// Note: there is no order guarantee for items entering/leaving the batcher.
func (b *basicBatcher) Put(item interface{}) error {
	b.lock.Lock()
	if b.disposed {
		b.lock.Unlock()
		return ErrDisposed
	}

	b.items = append(b.items, item)
	if b.calculateBytes != nil {
		b.availableBytes += b.calculateBytes(item)
	}
	if b.ready() {
		b.flush()
	}

	b.lock.Unlock()
	return nil
}

// Get retrieves a batch from the batcher. This call will block until
// one of the conditions for a "complete" batch is reached. If Put is
// continually called without calls to Get, an unbounded number of
// go-routines will be generated.
// Note: there is no order guarantee for items entering/leaving the batcher.
func (b *basicBatcher) Get() ([]interface{}, error) {
	// Don't check disposed yet so any items remaining in the queue
	// will be returned properly.

	var timeout <-chan time.Time
	if b.maxTime > 0 {
		timeout = time.After(b.maxTime)
	}

	select {
	case items, ok := <-b.batchChan:
		if !ok {
			return nil, ErrDisposed
		}
		return items, nil
	case <-timeout:
		b.lock.Lock()
		if b.disposed {
			b.lock.Unlock()
			return nil, ErrDisposed
		}
		items := b.items
		b.items = make([]interface{}, 0, b.maxItems)
		b.availableBytes = 0
		b.lock.Unlock()
		return items, nil
	}
}

// Flush forcibly completes the batch currently being built
func (b *basicBatcher) Flush() error {
	b.lock.Lock()
	if b.disposed {
		b.lock.Unlock()
		return ErrDisposed
	}
	b.flush()
	b.lock.Unlock()
	return nil
}

// Dispose will dispose of the batcher. Any calls to Put or Flush
// will return ErrDisposed, calls to Get will return an error iff
// there are no more ready batches.
func (b *basicBatcher) Dispose() {
	b.lock.Lock()
	if b.disposed {
		b.lock.Unlock()
		return
	}
	b.flush()
	b.disposed = true
	b.items = nil

	// Drain the batch channel and all routines waiting to put on the channel
	for len(b.batchChan) > 0 || atomic.LoadInt32(&b.waiting) > 0 {
		<-b.batchChan
	}
	close(b.batchChan)
	b.lock.Unlock()
}

// IsDisposed will determine if the batcher is disposed
func (b *basicBatcher) IsDisposed() bool {
	b.lock.RLock()
	disposed := b.disposed
	b.lock.RUnlock()
	return disposed
}

// flush adds the batch currently being built to the queue of completed batches.
// flush is not threadsafe, so should be synchronized externally.
func (b *basicBatcher) flush() {
	// Note: This needs to be in a go-routine to avoid locking out gets when
	// the batch channel is full.
	cpItems := make([]interface{}, len(b.items))
	for i, val := range b.items {
		cpItems[i] = val
	}
	// Signal one more waiter for the batch channel
	atomic.AddInt32(&b.waiting, 1)
	// Don't block on the channel put
	go func() {
		b.batchChan <- cpItems
		atomic.AddInt32(&b.waiting, -1)
	}()
	b.items = make([]interface{}, 0, b.maxItems)
	b.availableBytes = 0
}

func (b *basicBatcher) ready() bool {
	if b.maxItems != 0 && uint(len(b.items)) >= b.maxItems {
		return true
	}
	if b.maxBytes != 0 && b.availableBytes >= b.maxBytes {
		return true
	}
	return false
}
<commit_msg>Use an atomic pointer to track if batcher is disposed<commit_after>/*
Copyright 2015 Workiva, LLC

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

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

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

package batcher

import (
	"errors"
	"sync"
	"sync/atomic"
	"time"
)

const (
	batcherActive   = uint32(0)
	batcherDisposed = uint32(1)
)

// Batcher provides an API for accumulating items into a batch for processing.
type Batcher interface {
	// Put adds items to the batcher.
	Put(interface{}) error

	// Get retrieves a batch from the batcher. This call will block until
	// one of the conditions for a "complete" batch is reached.
	Get() ([]interface{}, error)

	// Flush forcibly completes the batch currently being built
	Flush() error

	// Dispose will dispose of the batcher. Any calls to Put or Flush
	// will return ErrDisposed, calls to Get will return an error iff
	// there are no more ready batches.
	Dispose()

	// IsDisposed will determine if the batcher is disposed
	IsDisposed() bool
}

// ErrDisposed is the error returned for a disposed Batcher
var ErrDisposed = errors.New("batcher: disposed")

// CalculateBytes evaluates the number of bytes in an item added to a Batcher.
type CalculateBytes func(interface{}) uint

type basicBatcher struct {
	maxTime        time.Duration
	maxItems       uint
	maxBytes       uint
	calculateBytes CalculateBytes
	disposed       uint32
	items          []interface{}
	lock           sync.RWMutex
	batchChan      chan []interface{}
	disposeChan    chan struct{}
	availableBytes uint
	waiting        int32
}

// New creates a new Batcher using the provided arguments.
// Batch readiness can be determined in three ways:
//   - Maximum number of bytes per batch
//   - Maximum number of items per batch
//   - Maximum amount of time waiting for a batch
// Values of zero for one of these fields indicate they should not be
// taken into account when evaluating the readiness of a batch.
func New(maxTime time.Duration, maxItems, maxBytes, queueLen uint, calculate CalculateBytes) (Batcher, error) {
	if maxBytes > 0 && calculate == nil {
		return nil, errors.New("batcher: must provide CalculateBytes function")
	}

	return &basicBatcher{
		maxTime:        maxTime,
		maxItems:       maxItems,
		maxBytes:       maxBytes,
		calculateBytes: calculate,
		items:          make([]interface{}, 0, maxItems),
		batchChan:      make(chan []interface{}, queueLen),
		disposeChan:    make(chan struct{}),
	}, nil
}

// Put adds items to the batcher. If Put is continually called without calls to
// Get, an unbounded number of go-routines will be generated.
// Note: there is no order guarantee for items entering/leaving the batcher.
func (b *basicBatcher) Put(item interface{}) error {
	// Check to see if disposed before putting
	if b.IsDisposed() {
		return ErrDisposed
	}
	b.lock.Lock()
	b.items = append(b.items, item)
	if b.calculateBytes != nil {
		b.availableBytes += b.calculateBytes(item)
	}
	if b.ready() {
		b.flush()
	}

	b.lock.Unlock()
	return nil
}

// Get retrieves a batch from the batcher. This call will block until
// one of the conditions for a "complete" batch is reached. If Put is
// continually called without calls to Get, an unbounded number of
// go-routines will be generated.
// Note: there is no order guarantee for items entering/leaving the batcher.
func (b *basicBatcher) Get() ([]interface{}, error) {
	// Don't check disposed yet so any items remaining in the queue
	// will be returned properly.

	var timeout <-chan time.Time
	if b.maxTime > 0 {
		timeout = time.After(b.maxTime)
	}

	// Check to see if disposed before blocking
	if b.IsDisposed() {
		return nil, ErrDisposed
	}

	select {
	case items := <-b.batchChan:
		return items, nil
	case _, ok := <-b.disposeChan:
		if !ok {
			return nil, ErrDisposed
		}
		return nil, nil
	case <-timeout:
		// Check to see if disposed before getting lock
		if b.IsDisposed() {
			return nil, ErrDisposed
		}
		b.lock.Lock()
		items := b.items
		b.items = make([]interface{}, 0, b.maxItems)
		b.availableBytes = 0
		b.lock.Unlock()
		return items, nil
	}
}

// Flush forcibly completes the batch currently being built
func (b *basicBatcher) Flush() error {
	if b.IsDisposed() {
		return ErrDisposed
	}
	b.lock.Lock()
	b.flush()
	b.lock.Unlock()
	return nil
}

// Dispose will dispose of the batcher. Any calls to Put or Flush
// will return ErrDisposed, calls to Get will return an error iff
// there are no more ready batches.
func (b *basicBatcher) Dispose() {
	// Check to see if disposed before attempting to dispose
	if atomic.CompareAndSwapUint32(&b.disposed, batcherActive, batcherDisposed) {
		return
	}
	b.lock.Lock()
	b.flush()
	b.items = nil
	close(b.disposeChan)

	// Drain the batch channel and all routines waiting to put on the channel
	for len(b.batchChan) > 0 || atomic.LoadInt32(&b.waiting) > 0 {
		<-b.batchChan
	}
	close(b.batchChan)
	b.lock.Unlock()
}

// IsDisposed will determine if the batcher is disposed
func (b *basicBatcher) IsDisposed() bool {
	return atomic.LoadUint32(&b.disposed) == batcherDisposed
}

// flush adds the batch currently being built to the queue of completed batches.
// flush is not threadsafe, so should be synchronized externally.
func (b *basicBatcher) flush() {
	// Note: This needs to be in a go-routine to avoid locking out gets when
	// the batch channel is full.
	cpItems := make([]interface{}, len(b.items))
	for i, val := range b.items {
		cpItems[i] = val
	}
	// Signal one more waiter for the batch channel
	atomic.AddInt32(&b.waiting, 1)
	// Don't block on the channel put
	go func() {
		b.batchChan <- cpItems
		atomic.AddInt32(&b.waiting, -1)
	}()
	b.items = make([]interface{}, 0, b.maxItems)
	b.availableBytes = 0
}

func (b *basicBatcher) ready() bool {
	if b.maxItems != 0 && uint(len(b.items)) >= b.maxItems {
		return true
	}
	if b.maxBytes != 0 && b.availableBytes >= b.maxBytes {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package commands

import (
	"errors"
	"fmt"
	"os"

	"github.com/arteev/uriban"

	"github.com/arteev/dsql/db"
	"github.com/arteev/dsql/parameters"
	"github.com/arteev/dsql/parameters/parametergetter"
	"github.com/arteev/dsql/rdb"
	"github.com/arteev/fmttab"
	"github.com/arteev/logger"
	"github.com/nsf/termbox-go"
	"github.com/urfave/cli"
)

func stringFlag(name, usage string) cli.Flag {
	return cli.StringFlag{
		Name:  name,
		Usage: usage,
	}
}

func listDatabase() cli.Command {

	dbFilterFlags := newCliFlags(cliOption{
		Databases:        modeFlagMulti,
		ExcludeDatabases: modeFlagMulti,
		Engines:          modeFlagMulti,
		Tags:             modeFlagMulti,
		ExcludeTags:      modeFlagMulti,
	})

	return cli.Command{
		Name:  "list",
		Usage: "list of databases",
		Flags: append(dbFilterFlags.Flags(),
			cli.BoolFlag{
				Name:  "fit",
				Usage: "use for fit table by width window of terminal",
			},
			cli.StringFlag{
				Name:  "border",
				Usage: "set type of border table: Thin,Double,Simple or None. Default:Thin",
			}),
		Action: func(ctx *cli.Context) error {
			logger.Trace.Println("command list database")
			dbFilterFlags.SetContext(ctx)
			d := db.GetInstance()
			dbs, err := d.All()
			if err != nil {
				return err
			}
			for _, e := range dbFilterFlags.Engines() {
				if err := rdb.CheckCodeEngine(e); err != nil {
					return err
				}
			}
			dbFilterFlags.ApplyTo(dbs)

			tab := fmttab.New("List of databases", fmttab.BorderThin, nil)
			tab.AddColumn("Id", 4, fmttab.AlignRight)
			tab.AddColumn("On", 2, fmttab.AlignLeft)
			tab.AddColumn("Code", 10, fmttab.AlignLeft)
			tab.AddColumn("Engine", 11, fmttab.AlignLeft)
			tab.AddColumn("URI", 40, fmttab.AlignLeft)
			tab.AddColumn("Tags", 25, fmttab.AlignLeft)
			for _, curd := range dbs.Get() {
				rec := make(map[string]interface{})
				rec["Id"] = curd.ID
				if curd.Enabled {
					rec["On"] = "+"
				}
				rec["Code"] = curd.Code
				rec["URI"] = uriban.Replace(curd.ConnectionString, uriban.WithOption(uriban.Password, uriban.ModeStarred(4)))
				rec["Engine"] = curd.Engine
				rec["Tags"] = curd.TagsComma(";")
				tab.AppendData(rec)
			}
			pget := parametergetter.New(ctx, parameters.GetInstance())
			if pget.GetDef(parametergetter.Fit, false).(bool) {
				if e := termbox.Init(); e != nil {
					panic(e)
				}
				tw, _ := termbox.Size()
				tab.AutoSize(true, tw)
				termbox.Close()
			}
			switch pget.GetDef(parametergetter.BorderTable, "").(string) {
			case "Thin":
				tab.SetBorder(fmttab.BorderThin)
			case "Double":
				tab.SetBorder(fmttab.BorderDouble)
			case "None":
				tab.SetBorder(fmttab.BorderNone)
			case "Simple":
				tab.SetBorder(fmttab.BorderSimple)
			}
			_, err = tab.WriteTo(os.Stdout)
			return err
		},
	}
}

func tagDatabase() cli.Command {
	dbFilterFlags := newCliFlags(cliOption{
		Databases:        modeFlagMulti,
		ExcludeDatabases: modeFlagMulti,
		Engines:          modeFlagMulti,
		Tags:             modeFlagUnUsed,
		ExcludeTags:      modeFlagUnUsed,
	})

	flags := dbFilterFlags.Flags()
	flags = append(flags, cli.StringSliceFlag{
		Name:  "add",
		Usage: "new tag(s)",
	})
	flags = append(flags, cli.StringSliceFlag{
		Name:  "remove",
		Usage: "remove tag(s)",
	})

	return cli.Command{
		Name:  "tag",
		Usage: "add or remove tag for database",
		Flags: flags,
		Action: func(ctx *cli.Context) error {
			logger.Trace.Println("command db tag")
			defer logger.Trace.Println("command db tag done")

			var add, remove = ctx.StringSlice("add"), ctx.StringSlice("remove")
			if len(add) == 0 && len(remove) == 0 {
				return fmt.Errorf("must be set: new tag or del tag")
			}
			dbFilterFlags.SetContext(ctx)

			logger.Debug.Printf("updating new:%s remove:%s\n", add, remove)

			d := db.GetInstance()
			col, err := d.All()
			if err != nil {
				return err
			}
			for _, e := range dbFilterFlags.Engines() {
				if err := rdb.CheckCodeEngine(e); err != nil {
					return err
				}
			}
			dbFilterFlags.ApplyTo(col)

			dbs := col.Get()
			if len(dbs) == 0 {
				return errors.New("databases not found")
			}

			pget := parametergetter.New(ctx, parameters.GetInstance())
			showstat := pget.GetDef(parametergetter.Statistic, false).(bool)
			var cntadd, cntremove int

			for _, curdb := range dbs {
				logger.Trace.Printf("process tag: %q\n", curdb.Code)

				cnt, err := d.AddTags(&curdb, add...)
				if err != nil {
					return err
				}
				cntadd += cnt
				logger.Info.Printf("Added tags %d for %s\n", cnt, curdb.Code)

				cnt, err = d.RemoveTags(&curdb, remove...)
				if err != nil {
					return err
				}
				cntremove += cnt
				logger.Info.Printf("Removed tags %d for %s\n", cnt, curdb.Code)

			}
			if showstat {
				fmt.Printf("Added tags: %d\nRemoved tags: %d\n", cntadd, cntremove)
			}

			return nil
		},
	}
}
func addDatabase() cli.Command {
	return cli.Command{
		Name:  "add",
		Usage: "Add new database",
		Flags: []cli.Flag{
			stringFlag("code", ""),
			stringFlag("uri", ""),
			stringFlag("engine", ""),
		},
		Action: func(ctx *cli.Context) error {
			logger.Trace.Println("command db add")
			for _, flag := range ctx.FlagNames() {
				if !ctx.IsSet(flag) {
					return fmt.Errorf("option %q must be set", flag)
				}
			}
			d := db.GetInstance()
			engine := ctx.String("engine")
			if err := rdb.CheckCodeEngine(engine); err != nil {
				return err
			}
			newdb := db.Database{
				Code:             ctx.String("code"),
				ConnectionString: ctx.String("uri"),
				Enabled:          true,
				Engine:           engine,
			}
			logger.Debug.Println("Adding ", newdb.Code, newdb.ConnectionString)
			err := d.Add(newdb)
			if err != nil {
				return err
			}
			logger.Info.Println("Added ", newdb.Code)
			return nil
		},
	}

}

func updateDatabase() cli.Command {
	return cli.Command{
		Name:  "update",
		Usage: "Update database",
		Flags: []cli.Flag{
			stringFlag("code", ""),
			stringFlag("newcode", ""),
			stringFlag("uri", ""),
			stringFlag("engine", ""),
			cli.BoolFlag{
				Name:  "enabled",
				Usage: "enabled or disable database",
			},
		},
		Action: func(ctx *cli.Context) error {
			logger.Trace.Println("command db update")
			if !ctx.IsSet("code") {
				return fmt.Errorf("option code must be set")
			}
			code := ctx.String("code")
			logger.Debug.Printf("updating %s, new values(code:%s; uri:%s; enabled:%v; engine:%v)\n", code, ctx.String("code"), ctx.String("uri"), ctx.Bool("enabled"), ctx.String("engine"))
			d := db.GetInstance()
			dbFind, err := d.FindByCode(code)
			logger.Debug.Println(dbFind)
			if err != nil {
				return err
			}
			if ctx.IsSet("newcode") {
				dbFind.Code = ctx.String("newcode")
			}
			if ctx.IsSet("uri") {
				dbFind.ConnectionString = ctx.String("uri")
			}
			if ctx.IsSet("enabled") {
				dbFind.Enabled = ctx.Bool("enabled")
			}
			if ctx.IsSet("engine") {
				dbFind.Engine = ctx.String("engine")
				if err := rdb.CheckCodeEngine(dbFind.Engine); err != nil {
					return err
				}
			}
			if err := d.Update(dbFind); err != nil {
				return err
			}

			logger.Info.Println("updated ", code)
			return nil
		},
	}

}

func deleteDatabase() cli.Command {
	return cli.Command{
		Name:  "delete",
		Usage: "Delete database by code",
		Flags: []cli.Flag{
			stringFlag("code", ""),
		},
		Action: func(ctx *cli.Context) {
			logger.Trace.Println("command database delete")
			if !ctx.IsSet("code") {
				panic(fmt.Errorf("option code must be set"))
			}
			code := ctx.String("code")
			logger.Debug.Printf("database deleting %q\n", code)
			d := db.GetInstance()
			dbfind, err := d.FindByCode(code)
			if err != nil {
				panic(err)
			}
			if err := d.Delete(dbfind); err != nil {
				panic(err)
			}
			logger.Info.Printf("database %q deleted\n", code)
		},
	}
}

//GetCommandsListDB define cli command for DB
func GetCommandsListDB() []cli.Command {
	return []cli.Command{
		cli.Command{
			Name:  "db",
			Usage: "list or managment of the list remote databases",
			Subcommands: []cli.Command{
				listDatabase(),
				addDatabase(),
				updateDatabase(),
				deleteDatabase(),
				tagDatabase(),
			},
		},
	}
}
<commit_msg>db list:autofit default<commit_after>package commands

import (
	"errors"
	"fmt"
	"os"

	"github.com/arteev/uriban"

	"github.com/arteev/dsql/db"
	"github.com/arteev/dsql/parameters"
	"github.com/arteev/dsql/parameters/parametergetter"
	"github.com/arteev/dsql/rdb"
	"github.com/arteev/fmttab"
	"github.com/arteev/logger"
	"github.com/nsf/termbox-go"
	"github.com/urfave/cli"
)

func stringFlag(name, usage string) cli.Flag {
	return cli.StringFlag{
		Name:  name,
		Usage: usage,
	}
}

func listDatabase() cli.Command {

	dbFilterFlags := newCliFlags(cliOption{
		Databases:        modeFlagMulti,
		ExcludeDatabases: modeFlagMulti,
		Engines:          modeFlagMulti,
		Tags:             modeFlagMulti,
		ExcludeTags:      modeFlagMulti,
	})

	return cli.Command{
		Name:  "list",
		Usage: "list of databases",
		Flags: append(dbFilterFlags.Flags(),
			cli.BoolFlag{
				Name:  "fit",
				Usage: "use for fit table by width window of terminal",
			},
			cli.StringFlag{
				Name:  "border",
				Usage: "set type of border table: Thin,Double,Simple or None. Default:Thin",
			}),
		Action: func(ctx *cli.Context) error {
			logger.Trace.Println("command list database")
			dbFilterFlags.SetContext(ctx)
			d := db.GetInstance()
			dbs, err := d.All()
			if err != nil {
				return err
			}
			for _, e := range dbFilterFlags.Engines() {
				if err := rdb.CheckCodeEngine(e); err != nil {
					return err
				}
			}
			dbFilterFlags.ApplyTo(dbs)

			tab := fmttab.New("List of databases", fmttab.BorderThin, nil)
			tab.AddColumn("Id", 4, fmttab.AlignRight)
			tab.AddColumn("On", 2, fmttab.AlignLeft)
			tab.AddColumn("Code", 10, fmttab.AlignLeft)
			tab.AddColumn("Engine", 11, fmttab.AlignLeft)
			tab.AddColumn("URI", 40, fmttab.AlignLeft)
			tab.AddColumn("Tags", 25, fmttab.AlignLeft)
			for _, curd := range dbs.Get() {
				rec := make(map[string]interface{})
				rec["Id"] = curd.ID
				if curd.Enabled {
					rec["On"] = "+"
				}
				rec["Code"] = curd.Code
				rec["URI"] = uriban.Replace(curd.ConnectionString, uriban.WithOption(uriban.Password, uriban.ModeStarred(4)))
				rec["Engine"] = curd.Engine
				rec["Tags"] = curd.TagsComma(";")
				tab.AppendData(rec)
			}
			pget := parametergetter.New(ctx, parameters.GetInstance())
			if pget.GetDef(parametergetter.Fit, true).(bool) {
				if err := termbox.Init(); err != nil {
					return err
				}
				tw, _ := termbox.Size()
				tab.AutoSize(true, tw)
				termbox.Close()
			}
			switch pget.GetDef(parametergetter.BorderTable, "").(string) {
			case "Thin":
				tab.SetBorder(fmttab.BorderThin)
			case "Double":
				tab.SetBorder(fmttab.BorderDouble)
			case "None":
				tab.SetBorder(fmttab.BorderNone)
			case "Simple":
				tab.SetBorder(fmttab.BorderSimple)
			}
			_, err = tab.WriteTo(os.Stdout)
			return err
		},
	}
}

func tagDatabase() cli.Command {
	dbFilterFlags := newCliFlags(cliOption{
		Databases:        modeFlagMulti,
		ExcludeDatabases: modeFlagMulti,
		Engines:          modeFlagMulti,
		Tags:             modeFlagUnUsed,
		ExcludeTags:      modeFlagUnUsed,
	})

	flags := dbFilterFlags.Flags()
	flags = append(flags, cli.StringSliceFlag{
		Name:  "add",
		Usage: "new tag(s)",
	})
	flags = append(flags, cli.StringSliceFlag{
		Name:  "remove",
		Usage: "remove tag(s)",
	})

	return cli.Command{
		Name:  "tag",
		Usage: "add or remove tag for database",
		Flags: flags,
		Action: func(ctx *cli.Context) error {
			logger.Trace.Println("command db tag")
			defer logger.Trace.Println("command db tag done")

			var add, remove = ctx.StringSlice("add"), ctx.StringSlice("remove")
			if len(add) == 0 && len(remove) == 0 {
				return fmt.Errorf("must be set: new tag or del tag")
			}
			dbFilterFlags.SetContext(ctx)

			logger.Debug.Printf("updating new:%s remove:%s\n", add, remove)

			d := db.GetInstance()
			col, err := d.All()
			if err != nil {
				return err
			}
			for _, e := range dbFilterFlags.Engines() {
				if err := rdb.CheckCodeEngine(e); err != nil {
					return err
				}
			}
			dbFilterFlags.ApplyTo(col)

			dbs := col.Get()
			if len(dbs) == 0 {
				return errors.New("databases not found")
			}

			pget := parametergetter.New(ctx, parameters.GetInstance())
			showstat := pget.GetDef(parametergetter.Statistic, false).(bool)
			var cntadd, cntremove int

			for _, curdb := range dbs {
				logger.Trace.Printf("process tag: %q\n", curdb.Code)

				cnt, err := d.AddTags(&curdb, add...)
				if err != nil {
					return err
				}
				cntadd += cnt
				logger.Info.Printf("Added tags %d for %s\n", cnt, curdb.Code)

				cnt, err = d.RemoveTags(&curdb, remove...)
				if err != nil {
					return err
				}
				cntremove += cnt
				logger.Info.Printf("Removed tags %d for %s\n", cnt, curdb.Code)

			}
			if showstat {
				fmt.Printf("Added tags: %d\nRemoved tags: %d\n", cntadd, cntremove)
			}

			return nil
		},
	}
}
func addDatabase() cli.Command {
	return cli.Command{
		Name:  "add",
		Usage: "Add new database",
		Flags: []cli.Flag{
			stringFlag("code", ""),
			stringFlag("uri", ""),
			stringFlag("engine", ""),
		},
		Action: func(ctx *cli.Context) error {
			logger.Trace.Println("command db add")
			for _, flag := range ctx.FlagNames() {
				if !ctx.IsSet(flag) {
					return fmt.Errorf("option %q must be set", flag)
				}
			}
			d := db.GetInstance()
			engine := ctx.String("engine")
			if err := rdb.CheckCodeEngine(engine); err != nil {
				return err
			}
			newdb := db.Database{
				Code:             ctx.String("code"),
				ConnectionString: ctx.String("uri"),
				Enabled:          true,
				Engine:           engine,
			}
			logger.Debug.Println("Adding ", newdb.Code, newdb.ConnectionString)
			err := d.Add(newdb)
			if err != nil {
				return err
			}
			logger.Info.Println("Added ", newdb.Code)
			return nil
		},
	}

}

func updateDatabase() cli.Command {
	return cli.Command{
		Name:  "update",
		Usage: "Update database",
		Flags: []cli.Flag{
			stringFlag("code", ""),
			stringFlag("newcode", ""),
			stringFlag("uri", ""),
			stringFlag("engine", ""),
			cli.BoolFlag{
				Name:  "enabled",
				Usage: "enabled or disable database",
			},
		},
		Action: func(ctx *cli.Context) error {
			logger.Trace.Println("command db update")
			if !ctx.IsSet("code") {
				return fmt.Errorf("option code must be set")
			}
			code := ctx.String("code")
			logger.Debug.Printf("updating %s, new values(code:%s; uri:%s; enabled:%v; engine:%v)\n", code, ctx.String("code"), ctx.String("uri"), ctx.Bool("enabled"), ctx.String("engine"))
			d := db.GetInstance()
			dbFind, err := d.FindByCode(code)
			logger.Debug.Println(dbFind)
			if err != nil {
				return err
			}
			if ctx.IsSet("newcode") {
				dbFind.Code = ctx.String("newcode")
			}
			if ctx.IsSet("uri") {
				dbFind.ConnectionString = ctx.String("uri")
			}
			if ctx.IsSet("enabled") {
				dbFind.Enabled = ctx.Bool("enabled")
			}
			if ctx.IsSet("engine") {
				dbFind.Engine = ctx.String("engine")
				if err := rdb.CheckCodeEngine(dbFind.Engine); err != nil {
					return err
				}
			}
			if err := d.Update(dbFind); err != nil {
				return err
			}

			logger.Info.Println("updated ", code)
			return nil
		},
	}

}

func deleteDatabase() cli.Command {
	return cli.Command{
		Name:  "delete",
		Usage: "Delete database by code",
		Flags: []cli.Flag{
			stringFlag("code", ""),
		},
		Action: func(ctx *cli.Context) {
			logger.Trace.Println("command database delete")
			if !ctx.IsSet("code") {
				panic(fmt.Errorf("option code must be set"))
			}
			code := ctx.String("code")
			logger.Debug.Printf("database deleting %q\n", code)
			d := db.GetInstance()
			dbfind, err := d.FindByCode(code)
			if err != nil {
				panic(err)
			}
			if err := d.Delete(dbfind); err != nil {
				panic(err)
			}
			logger.Info.Printf("database %q deleted\n", code)
		},
	}
}

//GetCommandsListDB define cli command for DB
func GetCommandsListDB() []cli.Command {
	return []cli.Command{
		cli.Command{
			Name:  "db",
			Usage: "list or managment of the list remote databases",
			Subcommands: []cli.Command{
				listDatabase(),
				addDatabase(),
				updateDatabase(),
				deleteDatabase(),
				tagDatabase(),
			},
		},
	}
}
<|endoftext|>


<commit_before>package actors

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"time"

	acceptance "github.com/cloudfoundry/bosh-bootloader/acceptance-tests"
	"github.com/kr/pty"

	"github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
)

type BBL struct {
	stateDirectory string
	pathToBBL      string
	configuration  acceptance.Config
	envID          string
}

func NewBBL(stateDirectory string, pathToBBL string, configuration acceptance.Config, envIDSuffix string) BBL {
	envIDPrefix := os.Getenv("BBL_TEST_ENV_ID_PREFIX")
	if envIDPrefix == "" {
		envIDPrefix = "bbl-test"
	}

	return BBL{
		stateDirectory: stateDirectory,
		pathToBBL:      pathToBBL,
		configuration:  configuration,
		envID:          fmt.Sprintf("%s-%s", envIDPrefix, envIDSuffix),
	}
}

func (b BBL) PredefinedEnvID() string {
	return b.envID
}

func (b BBL) Up(additionalArgs ...string) *gexec.Session {
	args := []string{
		"--state-dir", b.stateDirectory,
		"--debug",
		"up",
	}

	args = append(args, additionalArgs...)

	return b.execute(args, os.Stdout, os.Stderr)
}

func (b BBL) Plan(additionalArgs ...string) *gexec.Session {
	args := []string{
		"--state-dir", b.stateDirectory,
		"--debug",
		"plan",
	}

	args = append(args, additionalArgs...)

	return b.execute(args, os.Stdout, os.Stderr)
}

func (b BBL) Rotate() *gexec.Session {
	return b.execute([]string{
		"--state-dir", b.stateDirectory,
		"--debug",
		"rotate",
	}, os.Stdout, os.Stderr)
}

func (b BBL) VerifySSH(sshFunc func() (*exec.Cmd, *os.File)) {
	cmd, session := sshFunc()
	defer session.Close()

	fmt.Fprintln(session, "whoami")
	fmt.Fprintln(session, "exit 0")
	output, err := ioutil.ReadAll(session)
	Expect(err).NotTo(HaveOccurred())
	Expect(string(output)).To(ContainSubstring("jumpbox"))

	Eventually(cmd.Wait).Should(Succeed(), fmt.Sprintf("output was:\n\n%s", output))
}

func (b BBL) JumpboxSSH() (*exec.Cmd, *os.File) {
	return b.interactiveExecute([]string{
		"--state-dir", b.stateDirectory,
		"--debug",
		"ssh",
		"--jumpbox",
	}, os.Stdout, os.Stderr)
}

func (b BBL) DirectorSSH() (*exec.Cmd, *os.File) {
	return b.interactiveExecute([]string{
		"--state-dir", b.stateDirectory,
		"--debug",
		"ssh",
		"--director",
	}, os.Stdout, os.Stderr)
}

func (b BBL) Destroy() *gexec.Session {
	return b.execute([]string{
		"--state-dir", b.stateDirectory,
		"--debug",
		"destroy",
		"--no-confirm",
	}, os.Stdout, os.Stderr)
}

func (b BBL) Down() *gexec.Session {
	return b.execute([]string{
		"--state-dir", b.stateDirectory,
		"--debug",
		"down",
		"--no-confirm",
	}, os.Stdout, os.Stderr)
}

func (b BBL) CleanupLeftovers(filter string) *gexec.Session {
	return b.execute([]string{
		"--state-dir", b.stateDirectory,
		"cleanup-leftovers",
		"--filter", filter,
		"--no-confirm",
	}, os.Stdout, os.Stderr)
}

func (b BBL) Lbs() string {
	return b.fetchValue("lbs")
}

func (b BBL) DirectorUsername() string {
	return b.fetchValue("director-username")
}

func (b BBL) DirectorPassword() string {
	return b.fetchValue("director-password")
}

func (b BBL) DirectorAddress() string {
	return b.fetchValue("director-address")
}

func (b BBL) DirectorCACert() string {
	return b.fetchValue("director-ca-cert")
}

func (b BBL) JumpboxAddress() string {
	return b.fetchValue("jumpbox-address")
}

func (b BBL) SSHKey() string {
	return b.fetchValue("ssh-key")
}

func (b BBL) DirectorSSHKey() string {
	return b.fetchValue("director-ssh-key")
}

func (b BBL) EnvID() string {
	return b.fetchValue("env-id")
}

func (b BBL) PrintEnv() string {
	return b.fetchValue("print-env")
}

func (b BBL) LatestError() string {
	return b.fetchValue("latest-error")
}

func (b BBL) SaveDirectorCA() string {
	stdout := bytes.NewBuffer([]byte{})
	session := b.execute([]string{
		"--state-dir", b.stateDirectory,
		"director-ca-cert",
	}, stdout, os.Stderr)
	Eventually(session, 10*time.Minute).Should(gexec.Exit(0))

	file, err := ioutil.TempFile("", "")
	defer file.Close()
	Expect(err).NotTo(HaveOccurred())

	file.Write(stdout.Bytes())

	return file.Name()
}

func (b BBL) ExportBoshAllProxy() string {
	lines := strings.Split(b.PrintEnv(), "\n")
	value := getExport("BOSH_ALL_PROXY", lines)
	os.Setenv("BOSH_ALL_PROXY", value)
	return value
}

func (b BBL) StartSSHTunnel() *gexec.Session {
	printEnvLines := strings.Split(b.PrintEnv(), "\n")
	os.Setenv("BOSH_ALL_PROXY", getExport("BOSH_ALL_PROXY", printEnvLines))

	var sshArgs []string
	for i := 0; i < len(printEnvLines); i++ {
		if strings.HasPrefix(printEnvLines[i], "ssh ") {
			sshCmd := strings.TrimPrefix(printEnvLines[i], "ssh ")
			sshCmd = strings.Replace(sshCmd, "$JUMPBOX_PRIVATE_KEY", getExport("JUMPBOX_PRIVATE_KEY", printEnvLines), -1)
			sshCmd = strings.Replace(sshCmd, "-f ", "", -1)
			sshArgs = strings.Split(sshCmd, " ")
		}
	}

	cmd := exec.Command("ssh", sshArgs...)
	sshSession, err := gexec.Start(cmd, ginkgo.GinkgoWriter, ginkgo.GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())

	return sshSession
}

func getExport(keyName string, lines []string) string {
	for _, line := range lines {
		if strings.HasPrefix(line, fmt.Sprintf("export %s", keyName)) {
			parts := strings.Split(line, " ")
			keyValue := parts[1]
			keyValueParts := strings.Split(keyValue, "=")
			return strings.Join(keyValueParts[1:], "=")
		}
	}
	return ""
}

func (b BBL) fetchValue(value string) string {
	args := []string{
		"--state-dir", b.stateDirectory,
		value,
	}

	stdout := bytes.NewBuffer([]byte{})
	stderr := bytes.NewBuffer([]byte{})
	b.execute(args, stdout, stderr).Wait(30 * time.Second)

	return strings.TrimSpace(string(stdout.Bytes()))
}

func (b BBL) execute(args []string, stdout io.Writer, stderr io.Writer) *gexec.Session {
	cmd := exec.Command(b.pathToBBL, args...)
	session, err := gexec.Start(cmd, stdout, stderr)
	Expect(err).NotTo(HaveOccurred())

	return session
}

func (b BBL) interactiveExecute(args []string, stdout io.Writer, stderr io.Writer) (*exec.Cmd, *os.File) {
	cmd := exec.Command(b.pathToBBL, args...)
	f, err := pty.Start(cmd)
	Expect(err).NotTo(HaveOccurred())

	return cmd, f
}
<commit_msg>add sleeps to pty code because 'read /dev/ptmx: input/output error' is inscrutable<commit_after>package actors

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"time"

	acceptance "github.com/cloudfoundry/bosh-bootloader/acceptance-tests"
	"github.com/kr/pty"

	"github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
)

type BBL struct {
	stateDirectory string
	pathToBBL      string
	configuration  acceptance.Config
	envID          string
}

func NewBBL(stateDirectory string, pathToBBL string, configuration acceptance.Config, envIDSuffix string) BBL {
	envIDPrefix := os.Getenv("BBL_TEST_ENV_ID_PREFIX")
	if envIDPrefix == "" {
		envIDPrefix = "bbl-test"
	}

	return BBL{
		stateDirectory: stateDirectory,
		pathToBBL:      pathToBBL,
		configuration:  configuration,
		envID:          fmt.Sprintf("%s-%s", envIDPrefix, envIDSuffix),
	}
}

func (b BBL) PredefinedEnvID() string {
	return b.envID
}

func (b BBL) Up(additionalArgs ...string) *gexec.Session {
	args := []string{
		"--state-dir", b.stateDirectory,
		"--debug",
		"up",
	}

	args = append(args, additionalArgs...)

	return b.execute(args, os.Stdout, os.Stderr)
}

func (b BBL) Plan(additionalArgs ...string) *gexec.Session {
	args := []string{
		"--state-dir", b.stateDirectory,
		"--debug",
		"plan",
	}

	args = append(args, additionalArgs...)

	return b.execute(args, os.Stdout, os.Stderr)
}

func (b BBL) Rotate() *gexec.Session {
	return b.execute([]string{
		"--state-dir", b.stateDirectory,
		"--debug",
		"rotate",
	}, os.Stdout, os.Stderr)
}

func (b BBL) VerifySSH(sshFunc func() (*exec.Cmd, *os.File)) {
	cmd, session := sshFunc()
	defer session.Close()

	time.Sleep(5 * time.Second)
	fmt.Fprintln(session, "whoami")
	fmt.Fprintln(session, "exit 0")
	time.Sleep(5 * time.Second)
	output, err := ioutil.ReadAll(session)
	Expect(err).NotTo(HaveOccurred())
	Expect(string(output)).To(ContainSubstring("jumpbox"))

	Eventually(cmd.Wait).Should(Succeed(), fmt.Sprintf("output was:\n\n%s", output))
}

func (b BBL) JumpboxSSH() (*exec.Cmd, *os.File) {
	return b.interactiveExecute([]string{
		"--state-dir", b.stateDirectory,
		"--debug",
		"ssh",
		"--jumpbox",
	}, os.Stdout, os.Stderr)
}

func (b BBL) DirectorSSH() (*exec.Cmd, *os.File) {
	return b.interactiveExecute([]string{
		"--state-dir", b.stateDirectory,
		"--debug",
		"ssh",
		"--director",
	}, os.Stdout, os.Stderr)
}

func (b BBL) Destroy() *gexec.Session {
	return b.execute([]string{
		"--state-dir", b.stateDirectory,
		"--debug",
		"destroy",
		"--no-confirm",
	}, os.Stdout, os.Stderr)
}

func (b BBL) Down() *gexec.Session {
	return b.execute([]string{
		"--state-dir", b.stateDirectory,
		"--debug",
		"down",
		"--no-confirm",
	}, os.Stdout, os.Stderr)
}

func (b BBL) CleanupLeftovers(filter string) *gexec.Session {
	return b.execute([]string{
		"--state-dir", b.stateDirectory,
		"cleanup-leftovers",
		"--filter", filter,
		"--no-confirm",
	}, os.Stdout, os.Stderr)
}

func (b BBL) Lbs() string {
	return b.fetchValue("lbs")
}

func (b BBL) DirectorUsername() string {
	return b.fetchValue("director-username")
}

func (b BBL) DirectorPassword() string {
	return b.fetchValue("director-password")
}

func (b BBL) DirectorAddress() string {
	return b.fetchValue("director-address")
}

func (b BBL) DirectorCACert() string {
	return b.fetchValue("director-ca-cert")
}

func (b BBL) JumpboxAddress() string {
	return b.fetchValue("jumpbox-address")
}

func (b BBL) SSHKey() string {
	return b.fetchValue("ssh-key")
}

func (b BBL) DirectorSSHKey() string {
	return b.fetchValue("director-ssh-key")
}

func (b BBL) EnvID() string {
	return b.fetchValue("env-id")
}

func (b BBL) PrintEnv() string {
	return b.fetchValue("print-env")
}

func (b BBL) LatestError() string {
	return b.fetchValue("latest-error")
}

func (b BBL) SaveDirectorCA() string {
	stdout := bytes.NewBuffer([]byte{})
	session := b.execute([]string{
		"--state-dir", b.stateDirectory,
		"director-ca-cert",
	}, stdout, os.Stderr)
	Eventually(session, 10*time.Minute).Should(gexec.Exit(0))

	file, err := ioutil.TempFile("", "")
	defer file.Close()
	Expect(err).NotTo(HaveOccurred())

	file.Write(stdout.Bytes())

	return file.Name()
}

func (b BBL) ExportBoshAllProxy() string {
	lines := strings.Split(b.PrintEnv(), "\n")
	value := getExport("BOSH_ALL_PROXY", lines)
	os.Setenv("BOSH_ALL_PROXY", value)
	return value
}

func (b BBL) StartSSHTunnel() *gexec.Session {
	printEnvLines := strings.Split(b.PrintEnv(), "\n")
	os.Setenv("BOSH_ALL_PROXY", getExport("BOSH_ALL_PROXY", printEnvLines))

	var sshArgs []string
	for i := 0; i < len(printEnvLines); i++ {
		if strings.HasPrefix(printEnvLines[i], "ssh ") {
			sshCmd := strings.TrimPrefix(printEnvLines[i], "ssh ")
			sshCmd = strings.Replace(sshCmd, "$JUMPBOX_PRIVATE_KEY", getExport("JUMPBOX_PRIVATE_KEY", printEnvLines), -1)
			sshCmd = strings.Replace(sshCmd, "-f ", "", -1)
			sshArgs = strings.Split(sshCmd, " ")
		}
	}

	cmd := exec.Command("ssh", sshArgs...)
	sshSession, err := gexec.Start(cmd, ginkgo.GinkgoWriter, ginkgo.GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())

	return sshSession
}

func getExport(keyName string, lines []string) string {
	for _, line := range lines {
		if strings.HasPrefix(line, fmt.Sprintf("export %s", keyName)) {
			parts := strings.Split(line, " ")
			keyValue := parts[1]
			keyValueParts := strings.Split(keyValue, "=")
			return strings.Join(keyValueParts[1:], "=")
		}
	}
	return ""
}

func (b BBL) fetchValue(value string) string {
	args := []string{
		"--state-dir", b.stateDirectory,
		value,
	}

	stdout := bytes.NewBuffer([]byte{})
	stderr := bytes.NewBuffer([]byte{})
	b.execute(args, stdout, stderr).Wait(30 * time.Second)

	return strings.TrimSpace(string(stdout.Bytes()))
}

func (b BBL) execute(args []string, stdout io.Writer, stderr io.Writer) *gexec.Session {
	cmd := exec.Command(b.pathToBBL, args...)
	session, err := gexec.Start(cmd, stdout, stderr)
	Expect(err).NotTo(HaveOccurred())

	return session
}

func (b BBL) interactiveExecute(args []string, stdout io.Writer, stderr io.Writer) (*exec.Cmd, *os.File) {
	cmd := exec.Command(b.pathToBBL, args...)
	f, err := pty.Start(cmd)
	Expect(err).NotTo(HaveOccurred())

	return cmd, f
}
<|endoftext|>


<commit_before>package api

import (
	"crypto/tls"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"time"

	"github.com/grafana/grafana/pkg/api/cloudwatch"
	"github.com/grafana/grafana/pkg/bus"
	"github.com/grafana/grafana/pkg/middleware"
	m "github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/setting"
	"github.com/grafana/grafana/pkg/util"
)

var dataProxyTransport = &http.Transport{
	TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	Proxy:           http.ProxyFromEnvironment,
	Dial: (&net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
	}).Dial,
	TLSHandshakeTimeout: 10 * time.Second,
}

func NewReverseProxy(ds *m.DataSource, proxyPath string, targetUrl *url.URL) *httputil.ReverseProxy {
	director := func(req *http.Request) {
		req.URL.Scheme = targetUrl.Scheme
		req.URL.Host = targetUrl.Host
		req.Host = targetUrl.Host

		reqQueryVals := req.URL.Query()

		if ds.Type == m.DS_INFLUXDB_08 {
			req.URL.Path = util.JoinUrlFragments(targetUrl.Path, "db/"+ds.Database+"/"+proxyPath)
			reqQueryVals.Add("u", ds.User)
			reqQueryVals.Add("p", ds.Password)
			req.URL.RawQuery = reqQueryVals.Encode()
		} else if ds.Type == m.DS_INFLUXDB {
			req.URL.Path = util.JoinUrlFragments(targetUrl.Path, proxyPath)
			reqQueryVals.Add("db", ds.Database)
			req.URL.RawQuery = reqQueryVals.Encode()
			if !ds.BasicAuth {
				req.Header.Del("Authorization")
				req.Header.Add("Authorization", util.GetBasicAuthHeader(ds.User, ds.Password))
			}
		} else {
			req.URL.Path = util.JoinUrlFragments(targetUrl.Path, proxyPath)
		}

		if ds.BasicAuth {
			req.Header.Del("Authorization")
			req.Header.Add("Authorization", util.GetBasicAuthHeader(ds.BasicAuthUser, ds.BasicAuthPassword))
		}

		// clear cookie headers
		req.Header.Del("Cookie")
		req.Header.Del("Set-Cookie")
	}

	return &httputil.ReverseProxy{Director: director}
}

var dsMap map[int64]*m.DataSource = make(map[int64]*m.DataSource)

func getDatasource(id int64, orgId int64) (*m.DataSource, error) {
	// ds, exists := dsMap[id]
	// if exists && ds.OrgId == orgId {
	// 	return ds, nil
	// }

	query := m.GetDataSourceByIdQuery{Id: id, OrgId: orgId}
	if err := bus.Dispatch(&query); err != nil {
		return nil, err
	}

	dsMap[id] = &query.Result
	return &query.Result, nil
}

func ProxyDataSourceRequest(c *middleware.Context) {
	ds, err := getDatasource(c.ParamsInt64(":id"), c.OrgId)
	if err != nil {
		c.JsonApiErr(500, "Unable to load datasource meta data", err)
		return
	}

	targetUrl, _ := url.Parse(ds.Url)
	if len(setting.DataProxyWhiteList) > 0 {
		if _, exists := setting.DataProxyWhiteList[targetUrl.Host]; !exists {
			c.JsonApiErr(403, "Data proxy hostname and ip are not included in whitelist", nil)
			return
		}
	}

	if ds.Type == m.DS_CLOUDWATCH {
		cloudwatch.HandleRequest(c, ds)
	} else {
		proxyPath := c.Params("*")
		proxy := NewReverseProxy(ds, proxyPath, targetUrl)
		proxy.Transport = dataProxyTransport
		proxy.ServeHTTP(c.Resp, c.Req.Request)
		c.Resp.Header().Del("Set-Cookie")
	}
}
<commit_msg>fix(dataproxy): remove partially used cache<commit_after>package api

import (
	"crypto/tls"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"time"

	"github.com/grafana/grafana/pkg/api/cloudwatch"
	"github.com/grafana/grafana/pkg/bus"
	"github.com/grafana/grafana/pkg/middleware"
	m "github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/setting"
	"github.com/grafana/grafana/pkg/util"
)

var dataProxyTransport = &http.Transport{
	TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	Proxy:           http.ProxyFromEnvironment,
	Dial: (&net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
	}).Dial,
	TLSHandshakeTimeout: 10 * time.Second,
}

func NewReverseProxy(ds *m.DataSource, proxyPath string, targetUrl *url.URL) *httputil.ReverseProxy {
	director := func(req *http.Request) {
		req.URL.Scheme = targetUrl.Scheme
		req.URL.Host = targetUrl.Host
		req.Host = targetUrl.Host

		reqQueryVals := req.URL.Query()

		if ds.Type == m.DS_INFLUXDB_08 {
			req.URL.Path = util.JoinUrlFragments(targetUrl.Path, "db/"+ds.Database+"/"+proxyPath)
			reqQueryVals.Add("u", ds.User)
			reqQueryVals.Add("p", ds.Password)
			req.URL.RawQuery = reqQueryVals.Encode()
		} else if ds.Type == m.DS_INFLUXDB {
			req.URL.Path = util.JoinUrlFragments(targetUrl.Path, proxyPath)
			reqQueryVals.Add("db", ds.Database)
			req.URL.RawQuery = reqQueryVals.Encode()
			if !ds.BasicAuth {
				req.Header.Del("Authorization")
				req.Header.Add("Authorization", util.GetBasicAuthHeader(ds.User, ds.Password))
			}
		} else {
			req.URL.Path = util.JoinUrlFragments(targetUrl.Path, proxyPath)
		}

		if ds.BasicAuth {
			req.Header.Del("Authorization")
			req.Header.Add("Authorization", util.GetBasicAuthHeader(ds.BasicAuthUser, ds.BasicAuthPassword))
		}

		// clear cookie headers
		req.Header.Del("Cookie")
		req.Header.Del("Set-Cookie")
	}

	return &httputil.ReverseProxy{Director: director}
}

func getDatasource(id int64, orgId int64) (*m.DataSource, error) {
	query := m.GetDataSourceByIdQuery{Id: id, OrgId: orgId}
	if err := bus.Dispatch(&query); err != nil {
		return nil, err
	}

	return &query.Result, nil
}

func ProxyDataSourceRequest(c *middleware.Context) {
	ds, err := getDatasource(c.ParamsInt64(":id"), c.OrgId)
	if err != nil {
		c.JsonApiErr(500, "Unable to load datasource meta data", err)
		return
	}

	targetUrl, _ := url.Parse(ds.Url)
	if len(setting.DataProxyWhiteList) > 0 {
		if _, exists := setting.DataProxyWhiteList[targetUrl.Host]; !exists {
			c.JsonApiErr(403, "Data proxy hostname and ip are not included in whitelist", nil)
			return
		}
	}

	if ds.Type == m.DS_CLOUDWATCH {
		cloudwatch.HandleRequest(c, ds)
	} else {
		proxyPath := c.Params("*")
		proxy := NewReverseProxy(ds, proxyPath, targetUrl)
		proxy.Transport = dataProxyTransport
		proxy.ServeHTTP(c.Resp, c.Req.Request)
		c.Resp.Header().Del("Set-Cookie")
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2018 the Heptio Ark contributors.

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 client

import (
	"encoding/json"
	"os"
	"path/filepath"

	"github.com/pkg/errors"
)

const (
	ConfigKeyNamespace = "namespace"
)

// LoadConfig loads the Ark client configuration file and returns it as a map[string]string. If the
// file does not exist, an empty map is returned.
func LoadConfig() (map[string]string, error) {
	fileName := configFileName()

	_, err := os.Stat(fileName)
	if os.IsNotExist(err) {
		// If the file isn't there, just return an empty map
		return map[string]string{}, nil
	}
	if err != nil {
		// For any other Stat() error, return it
		return nil, errors.WithStack(err)
	}

	configFile, err := os.Open(fileName)
	if err != nil {
		return nil, errors.WithStack(err)
	}
	defer configFile.Close()

	var config map[string]string
	if err := json.NewDecoder(configFile).Decode(&config); err != nil {
		return nil, errors.WithStack(err)
	}

	return config, nil
}

// SaveConfig saves the passed in config map to the Ark client configuration file.
func SaveConfig(config map[string]string) error {
	fileName := configFileName()

	configFile, err := os.OpenFile(fileName, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
	if err != nil {
		return errors.WithStack(err)
	}
	defer configFile.Close()

	return json.NewEncoder(configFile).Encode(&config)
}

func configFileName() string {
	return filepath.Join(os.Getenv("HOME"), ".config", "ark", "config.json")
}
<commit_msg>Create the config directory in case it's missing<commit_after>/*
Copyright 2018 the Heptio Ark contributors.

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 client

import (
	"encoding/json"
	"os"
	"path/filepath"

	"github.com/pkg/errors"
)

const (
	ConfigKeyNamespace = "namespace"
)

// LoadConfig loads the Ark client configuration file and returns it as a map[string]string. If the
// file does not exist, an empty map is returned.
func LoadConfig() (map[string]string, error) {
	fileName := configFileName()

	_, err := os.Stat(fileName)
	if os.IsNotExist(err) {
		// If the file isn't there, just return an empty map
		return map[string]string{}, nil
	}
	if err != nil {
		// For any other Stat() error, return it
		return nil, errors.WithStack(err)
	}

	configFile, err := os.Open(fileName)
	if err != nil {
		return nil, errors.WithStack(err)
	}
	defer configFile.Close()

	var config map[string]string
	if err := json.NewDecoder(configFile).Decode(&config); err != nil {
		return nil, errors.WithStack(err)
	}

	return config, nil
}

// SaveConfig saves the passed in config map to the Ark client configuration file.
func SaveConfig(config map[string]string) error {
	fileName := configFileName()

	// Try to make the directory in case it doesn't exist
	dir := filepath.Dir(fileName)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return errors.WithStack(err)
	}

	configFile, err := os.OpenFile(fileName, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
	if err != nil {
		return errors.WithStack(err)
	}
	defer configFile.Close()

	return json.NewEncoder(configFile).Encode(&config)
}

func configFileName() string {
	return filepath.Join(os.Getenv("HOME"), ".config", "ark", "config.json")
}
<|endoftext|>