<commit_before>// +-------------------------------------------------------------------------
// | Copyright (C) 2016 Yunify, Inc.
// +-------------------------------------------------------------------------
// | Licensed under the Apache License, Version 2.0 (the "License");
// | you may not use this work except in compliance with the License.
// | You may obtain a copy of the License in the LICENSE file, or 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 unpacker

import (
	"bytes"
	"encoding/json"
	"net/http"
	"reflect"

	"github.com/yunify/qingstor-sdk-go/request/data"
	"github.com/yunify/qingstor-sdk-go/request/errors"
)

// QingStorUnpacker is the response unpacker for QingStor service.
type QingStorUnpacker struct {
	baseUnpacker *BaseUnpacker
}

// UnpackHTTPRequest unpack the http response with an operation, http response and an output.
func (qu *QingStorUnpacker) UnpackHTTPRequest(o *data.Operation, r *http.Response, x *reflect.Value) error {
	qu.baseUnpacker = &BaseUnpacker{}
	err := qu.baseUnpacker.UnpackHTTPRequest(o, r, x)
	if err != nil {
		return err
	}

	err = qu.parseError()
	if err != nil {
		return err
	}

	return nil
}

func (qu *QingStorUnpacker) parseError() error {
	if !qu.baseUnpacker.isResponseRight() {
		if qu.baseUnpacker.httpResponse.Header.Get("Content-Type") == "application/json" {
			buffer := &bytes.Buffer{}
			buffer.ReadFrom(qu.baseUnpacker.httpResponse.Body)
			qu.baseUnpacker.httpResponse.Body.Close()

			qsError := &errors.QingStorError{}
			if buffer.Len() > 0 {
				err := json.Unmarshal(buffer.Bytes(), qsError)
				if err != nil {
					return err
				}
			}
			qsError.StatusCode = qu.baseUnpacker.httpResponse.StatusCode

			return qsError
		}
	}

	return nil
}
<commit_msg>Close body for every API except GetObject and ImageProcess<commit_after>// +-------------------------------------------------------------------------
// | Copyright (C) 2016 Yunify, Inc.
// +-------------------------------------------------------------------------
// | Licensed under the Apache License, Version 2.0 (the "License");
// | you may not use this work except in compliance with the License.
// | You may obtain a copy of the License in the LICENSE file, or 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 unpacker

import (
	"bytes"
	"encoding/json"
	"net/http"
	"reflect"

	"github.com/yunify/qingstor-sdk-go/request/data"
	"github.com/yunify/qingstor-sdk-go/request/errors"
)

// QingStorUnpacker is the response unpacker for QingStor service.
type QingStorUnpacker struct {
	baseUnpacker *BaseUnpacker
}

// UnpackHTTPRequest unpack the http response with an operation, http response and an output.
func (qu *QingStorUnpacker) UnpackHTTPRequest(o *data.Operation, r *http.Response, x *reflect.Value) error {
	qu.baseUnpacker = &BaseUnpacker{}
	err := qu.baseUnpacker.UnpackHTTPRequest(o, r, x)
	if err != nil {
		return err
	}

	err = qu.parseError()
	if err != nil {
		return err
	}

	// Close body for every API except GetObject and ImageProcess.
	if o.APIName != "GET Object" && o.APIName != "Image Process" && r.Body != nil {
		err = r.Body.Close()
		if err != nil {
			return err
		}

		r.Body = nil
	}

	return nil
}

func (qu *QingStorUnpacker) parseError() error {
	if !qu.baseUnpacker.isResponseRight() {
		if qu.baseUnpacker.httpResponse.Header.Get("Content-Type") == "application/json" {
			buffer := &bytes.Buffer{}
			buffer.ReadFrom(qu.baseUnpacker.httpResponse.Body)
			qu.baseUnpacker.httpResponse.Body.Close()

			qsError := &errors.QingStorError{}
			if buffer.Len() > 0 {
				err := json.Unmarshal(buffer.Bytes(), qsError)
				if err != nil {
					return err
				}
			}
			qsError.StatusCode = qu.baseUnpacker.httpResponse.StatusCode

			return qsError
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package requester

import (
	"errors"
	"strconv"
	"time"

	"github.com/Shopify/sarama"
	"github.com/tylertreat/bench"
)

// KafkaRequesterFactory implements RequesterFactory by creating a Requester
// which publishes messages to Kafka and waits to consume them.
type KafkaRequesterFactory struct {
	URLs        []string
	PayloadSize int
	Topic       string
}

// GetRequester returns a new Requester, called for each Benchmark connection.
func (k *KafkaRequesterFactory) GetRequester(num uint64) bench.Requester {
	return &kafkaRequester{
		urls:        k.URLs,
		payloadSize: k.PayloadSize,
		topic:       k.Topic + "-" + strconv.FormatUint(num, 10),
	}
}

// kafkaRequester implements Requester by publishing a message to Kafka and
// waiting to consume it.
type kafkaRequester struct {
	urls              []string
	payloadSize       int
	topic             string
	producer          sarama.SyncProducer
	consumer          sarama.Consumer
	partitionConsumer sarama.PartitionConsumer
	msg               *sarama.ProducerMessage
}

// Setup prepares the Requester for benchmarking.
func (k *kafkaRequester) Setup() error {
	config := sarama.NewConfig()
	producer, err := sarama.NewSyncProducer(k.urls, config)
	if err != nil {
		return err
	}

	consumer, err := sarama.NewConsumer(k.urls, nil)
	if err != nil {
		producer.Close()
		return err
	}
	partitionConsumer, err := consumer.ConsumePartition(k.topic, 0, sarama.OffsetNewest)
	if err != nil {
		producer.Close()
		consumer.Close()
		return err
	}

	k.producer = producer
	k.consumer = consumer
	k.partitionConsumer = partitionConsumer
	k.msg = &sarama.ProducerMessage{
		Topic: k.topic,
		Value: sarama.ByteEncoder(make([]byte, k.payloadSize)),
	}
	return nil
}

// Request performs a synchronous request to the system under test.
func (k *kafkaRequester) Request() error {
	_, _, err := k.producer.SendMessage(k.msg)
	if err != nil {
		return err
	}

	select {
	case <-k.partitionConsumer.Messages():
		return nil
	case <-time.After(30 * time.Second):
		return errors.New("requester: Request timed out receiving")
	}
}

// Teardown is called upon benchmark completion.
func (k *kafkaRequester) Teardown() error {
	if err := k.partitionConsumer.Close(); err != nil {
		return err
	}
	if err := k.consumer.Close(); err != nil {
		return err
	}
	if err := k.producer.Close(); err != nil {
		return err
	}
	k.partitionConsumer = nil
	k.consumer = nil
	k.producer = nil
	return nil
}
<commit_msg>Use AsyncProducer in Kafka Requester<commit_after>package requester

import (
	"errors"
	"strconv"
	"time"

	"github.com/Shopify/sarama"
	"github.com/tylertreat/bench"
)

// KafkaRequesterFactory implements RequesterFactory by creating a Requester
// which publishes messages to Kafka and waits to consume them.
type KafkaRequesterFactory struct {
	URLs        []string
	PayloadSize int
	Topic       string
}

// GetRequester returns a new Requester, called for each Benchmark connection.
func (k *KafkaRequesterFactory) GetRequester(num uint64) bench.Requester {
	return &kafkaRequester{
		urls:        k.URLs,
		payloadSize: k.PayloadSize,
		topic:       k.Topic + "-" + strconv.FormatUint(num, 10),
	}
}

// kafkaRequester implements Requester by publishing a message to Kafka and
// waiting to consume it.
type kafkaRequester struct {
	urls              []string
	payloadSize       int
	topic             string
	producer          sarama.AsyncProducer
	consumer          sarama.Consumer
	partitionConsumer sarama.PartitionConsumer
	msg               *sarama.ProducerMessage
}

// Setup prepares the Requester for benchmarking.
func (k *kafkaRequester) Setup() error {
	config := sarama.NewConfig()
	producer, err := sarama.NewAsyncProducer(k.urls, config)
	if err != nil {
		return err
	}

	consumer, err := sarama.NewConsumer(k.urls, nil)
	if err != nil {
		producer.Close()
		return err
	}
	partitionConsumer, err := consumer.ConsumePartition(k.topic, 0, sarama.OffsetNewest)
	if err != nil {
		producer.Close()
		consumer.Close()
		return err
	}

	k.producer = producer
	k.consumer = consumer
	k.partitionConsumer = partitionConsumer
	k.msg = &sarama.ProducerMessage{
		Topic: k.topic,
		Value: sarama.ByteEncoder(make([]byte, k.payloadSize)),
	}

	return nil
}

// Request performs a synchronous request to the system under test.
func (k *kafkaRequester) Request() error {
	k.producer.Input() <- k.msg
	select {
	case <-k.partitionConsumer.Messages():
		return nil
	case <-time.After(30 * time.Second):
		return errors.New("requester: Request timed out receiving")
	}
}

// Teardown is called upon benchmark completion.
func (k *kafkaRequester) Teardown() error {
	if err := k.partitionConsumer.Close(); err != nil {
		return err
	}
	if err := k.consumer.Close(); err != nil {
		return err
	}
	if err := k.producer.Close(); err != nil {
		return err
	}
	k.partitionConsumer = nil
	k.consumer = nil
	k.producer = nil
	return nil
}
<|endoftext|>


<commit_before>
package main

import (
  "time"
  "bytes"
  "net"
  "net/http"
  "os/exec"
  "errors"
  "strings"
  "fmt"
  "log"
  "encoding/json"
)

type Info interface {
  Encode() ([]byte, error);
}

type Module interface {
  Monitor() (Info, error);
}

type Reporter struct {
  config *Config;
  module Module;
  running bool;
  moduleType string;
  url string;
}

type HostInfo struct {
  Hostnames    []string  `json:"hostnames"`;
  IpAddresses  []string  `json:"ipAddresses"`;
  Uname        string    `json:"uname"`;
  UnameA       string    `json:"unameA"`;
}

func (this *Reporter) Report() (error) {
  info, err := this.module.Monitor();
  if err != nil {
    return err;
  }
  reportData, err := info.Encode();
  if err != nil {
    return err
  }
  hostData, err := json.Marshal(this.GetHostInfo());
  envelope := fmt.Sprintf(`{"host":%s,"report":%s}`, hostData, reportData);
  req, err := http.NewRequest("POST", this.url, bytes.NewBuffer([]byte(envelope)))
  req.Header.Set("Content-Type", "application/json")
  client := &http.Client{}
  resp, err := client.Do(req)
  if err != nil {
    return err
  }
  log.Printf("Reported %s status to %s\n", this.moduleType, this.url);
  defer resp.Body.Close()
  if resp.StatusCode != 200 {
    return errors.New("Invalid response from Atoll server: " + resp.Status)
  }
  return nil
}

func (this *Reporter) GetHostInfo() (HostInfo) {
  hostInfo := HostInfo{};
  hostInfo.Hostnames = make([]string, 0);
  if len(this.config.Hostname) > 0 {
    hostInfo.Hostnames = append(hostInfo.Hostnames, this.config.Hostname);
  }

  data, err := exec.Command("hostname").Output();
  if err == nil {
    hostInfo.Hostnames = append(hostInfo.Hostnames, strings.TrimSpace(string(data)));
  }

  data, err = exec.Command("uname").Output();
  if err == nil {
    hostInfo.Uname = strings.TrimSpace(string(data));
  }

  data, err = exec.Command("uname", "-a").Output();
  if err == nil {
    hostInfo.UnameA = strings.TrimSpace(string(data));
  }

  hostInfo.IpAddresses = this.GetIpAddresses();

  for _,ipAddress := range hostInfo.IpAddresses {
    hostnames, err := net.LookupAddr(ipAddress)
    if err == nil {
      for _,hostname := range hostnames {
        hostInfo.Hostnames = append(hostInfo.Hostnames, hostname);
      }
    }
  }

  return hostInfo;
}

func (this *Reporter) GetIpAddresses() ([]string) {
  ipAddresses := make([]string, 0);
  interfaces, err := net.Interfaces()
  if err != nil {
    log.Printf("Warning could not get net.Interfaces information: %s\n", err);
    return ipAddresses
  }
  for _, i := range interfaces {
    addrs, err := i.Addrs()
    if err == nil {
      for _, addr := range addrs {
        var ip net.IP
        switch v := addr.(type) {
          case *net.IPNet:
            ip = v.IP
          case *net.IPAddr:
            ip = v.IP
        }
        ipnet := addr.(*net.IPNet);
        if (!ipnet.IP.IsLoopback()) {
          ipAddresses = append(ipAddresses, ip.String());
        }
      }
    }
  }
  return ipAddresses;
}

func (this *Reporter) Start() (error) {
  numSeconds, err := IntervalToSeconds(this.config.Publish.Frequency);
  if numSeconds == 0 {
    return errors.New("Need an interval of at least 1 second")
  }
  if err != nil {
    return errors.New("Invalid interval configured")
  }
  for this.running == true {
    err := this.Report();
    if err != nil {
      log.Printf("Warning, could not report: %v\n", err);
    }
    time.Sleep(time.Duration(numSeconds) * time.Second);
  }
  return nil
}

func (this *Reporter) Stop() () {
  this.running = false;
}
<commit_msg>Making sure no duplicate hostnames are sent<commit_after>
package main

import (
  "time"
  "bytes"
  "net"
  "net/http"
  "os/exec"
  "errors"
  "strings"
  "fmt"
  "log"
  "encoding/json"
)

type Info interface {
  Encode() ([]byte, error);
}

type Module interface {
  Monitor() (Info, error);
}

type Reporter struct {
  config *Config;
  module Module;
  running bool;
  moduleType string;
  url string;
}

type HostInfo struct {
  Hostnames    []string  `json:"hostnames"`;
  IpAddresses  []string  `json:"ipAddresses"`;
  Uname        string    `json:"uname"`;
  UnameA       string    `json:"unameA"`;
}

func (this *Reporter) Report() (error) {
  info, err := this.module.Monitor();
  if err != nil {
    return err;
  }
  reportData, err := info.Encode();
  if err != nil {
    return err
  }
  hostData, err := json.Marshal(this.GetHostInfo());
  envelope := fmt.Sprintf(`{"host":%s,"report":%s}`, hostData, reportData);
  req, err := http.NewRequest("POST", this.url, bytes.NewBuffer([]byte(envelope)))
  req.Header.Set("Content-Type", "application/json")
  client := &http.Client{}
  resp, err := client.Do(req)
  if err != nil {
    return err
  }
  log.Printf("Reported %s status to %s\n", this.moduleType, this.url);
  defer resp.Body.Close()
  if resp.StatusCode != 200 {
    return errors.New("Invalid response from Atoll server: " + resp.Status)
  }
  return nil
}

func arrayContains(s []string, e string) bool {
    for _, a := range s {
        if a == e {
            return true
        }
    }
    return false
}

func (this *Reporter) GetHostInfo() (HostInfo) {
  hostInfo := HostInfo{};
  hostInfo.Hostnames = make([]string, 0);
  if len(this.config.Hostname) > 0 {
    hostInfo.Hostnames = append(hostInfo.Hostnames, this.config.Hostname);
  }

  data, err := exec.Command("hostname").Output();
  if err == nil {
    hostInfo.Hostnames = append(hostInfo.Hostnames, strings.TrimSpace(string(data)));
  }

  data, err = exec.Command("uname").Output();
  if err == nil {
    hostInfo.Uname = strings.TrimSpace(string(data));
  }

  data, err = exec.Command("uname", "-a").Output();
  if err == nil {
    hostInfo.UnameA = strings.TrimSpace(string(data));
  }

  hostInfo.IpAddresses = this.GetIpAddresses();

  for _,ipAddress := range hostInfo.IpAddresses {
    hostnames, err := net.LookupAddr(ipAddress)
    if err == nil {
      for _,hostname := range hostnames {
        if arrayContains(hostInfo.Hostnames, hostname) == false {
          hostInfo.Hostnames = append(hostInfo.Hostnames, hostname);
        }
      }
    }
  }

  return hostInfo;
}

func (this *Reporter) GetIpAddresses() ([]string) {
  ipAddresses := make([]string, 0);
  interfaces, err := net.Interfaces()
  if err != nil {
    log.Printf("Warning could not get net.Interfaces information: %s\n", err);
    return ipAddresses
  }
  for _, i := range interfaces {
    addrs, err := i.Addrs()
    if err == nil {
      for _, addr := range addrs {
        var ip net.IP
        switch v := addr.(type) {
          case *net.IPNet:
            ip = v.IP
          case *net.IPAddr:
            ip = v.IP
        }
        ipnet := addr.(*net.IPNet);
        if (!ipnet.IP.IsLoopback()) {
          ipAddresses = append(ipAddresses, ip.String());
        }
      }
    }
  }
  return ipAddresses;
}

func (this *Reporter) Start() (error) {
  numSeconds, err := IntervalToSeconds(this.config.Publish.Frequency);
  if numSeconds == 0 {
    return errors.New("Need an interval of at least 1 second")
  }
  if err != nil {
    return errors.New("Invalid interval configured")
  }
  for this.running == true {
    err := this.Report();
    if err != nil {
      log.Printf("Warning, could not report: %v\n", err);
    }
    time.Sleep(time.Duration(numSeconds) * time.Second);
  }
  return nil
}

func (this *Reporter) Stop() () {
  this.running = false;
}
<|endoftext|>


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

// Package repotool is able to fetch information from a source code repository.
// Typically, it can get all commits, their authors and commiters and so on
// and return this information in a JSON object. Alternatively, it is able
// to populate the information into a PostgreSQL database.
// Currently, on the Git VCS is supported.
package main

import (
	"bufio"
	"bytes"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	_ "github.com/lib/pq"
	mmh3 "github.com/spaolacci/murmur3"

	"github.com/DevMine/srcanlzr/src"

	"github.com/DevMine/repotool/config"
	"github.com/DevMine/repotool/model"
	"github.com/DevMine/repotool/repo"
)

const version = "0.1.0"

// database fields per tables
var (
	diffDeltaFields = []string{
		"commit_id",
		"file_status",
		"is_file_binary",
		"similarity",
		"old_file_path",
		"new_file_path"}

	commitFields = []string{
		"repository_id",
		"author_id",
		"committer_id",
		"hash",
		"vcs_id",
		"message",
		"author_date",
		"commit_date",
		"file_changed_count",
		"insertions_count",
		"deletions_count"}
)

func main() {
	flag.Usage = func() {
		fmt.Printf("usage: %s [OPTION(S)] [REPOSITORY PATH]\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(0)
	}

	configPath := flag.String("c", "", "configuration file")
	vflag := flag.Bool("v", false, "print version.")
	jsonflag := flag.Bool("json", true, "json output")
	dbflag := flag.Bool("db", false, "import data into the database")
	srctoolflag := flag.String("srctool", "", "read json file produced by srctool (give stdin to read from stdin)")
	flag.Parse()

	if *vflag {
		fmt.Printf("%s - %s\n", filepath.Base(os.Args[0]), version)
		os.Exit(0)
	}

	if len(flag.Args()) != 1 {
		fmt.Fprintln(os.Stderr, "invalid # of arguments")
		flag.Usage()
	}

	if *dbflag && len(*configPath) == 0 {
		fatal(errors.New("a configuration file must be specified when using db option"))
	}

	if !*jsonflag && (len(*srctoolflag) > 0) {
		fatal(errors.New("srctool flag may be used only in conjonction with json flag"))
	}

	cfg, err := config.ReadConfig(*configPath)
	if err != nil {
		fatal(err)
	}

	repoPath := flag.Arg(0)
	repository, err := repo.New(*cfg, repoPath)
	if err != nil {
		fatal(err)
	}
	defer repository.Cleanup()

	fmt.Fprintln(os.Stderr, "fetching repository commits...")
	tic := time.Now()
	err = repository.FetchCommits()
	if err != nil {
		fatal(err)
	}
	toc := time.Now()
	fmt.Fprintln(os.Stderr, "done in ", toc.Sub(tic))

	if *jsonflag && (len(*srctoolflag) == 0) {
		bs, err := json.Marshal(repository)
		if err != nil {
			fatal(err)
		}
		fmt.Println(string(bs))
	}

	if *jsonflag && (len(*srctoolflag)) > 0 {
		var r *bufio.Reader
		if *srctoolflag == strings.ToLower("stdin") {
			// read from stdin
			r = bufio.NewReader(os.Stdin)
		} else {
			// read from srctool json file
			var f *os.File
			if f, err = os.Open(*srctoolflag); err != nil {
				fatal(err)
			}
			r = bufio.NewReader(f)
		}

		buf := new(bytes.Buffer)
		if _, err = io.Copy(buf, r); err != nil {
			fatal(err)
		}

		bs := buf.Bytes()
		p, err := src.Unmarshal(bs)
		if err != nil {
			fatal(err)
		}

		p.Repo = repository.GetRepository()
		bs, err = src.Marshal(p)
		if err != nil {
			fatal(err)
		}

		fmt.Println(string(bs))
	}

	if *dbflag {
		db, err := openDBSession(cfg.Database)
		if err != nil {
			fatal(err)
		}
		defer db.Close()

		fmt.Fprintf(os.Stderr,
			"inserting %d commits from %s repository into the database...\n",
			len(repository.GetCommits()), repository.GetName())
		tic := time.Now()
		err = insertRepoData(db, repository)
		toc := time.Now()
		if err != nil {
			fmt.Fprintln(os.Stderr, err.Error())
			os.Exit(1)
		} else {
			fmt.Fprintln(os.Stderr, "done in ", toc.Sub(tic))
		}
	}
}

// fatal prints an error on standard error stream and exits.
func fatal(a ...interface{}) {
	fmt.Fprintln(os.Stderr, a...)
	os.Exit(1)
}

// fail prints an error on standard error stream.
func fail(a ...interface{}) {
	fmt.Fprintln(os.Stderr, a...)
}

// openDBSession creates a session to the database.
func openDBSession(cfg config.DatabaseConfig) (*sql.DB, error) {
	dbURL := fmt.Sprintf(
		"user='%s' password='%s' host='%s' port=%d dbname='%s' sslmode='%s'",
		cfg.UserName, cfg.Password, cfg.HostName, cfg.Port, cfg.DBName, cfg.SSLMode)

	return sql.Open("postgres", dbURL)
}

// insertRepoData inserts repository data into the database, or updates it
// if it is already there.
func insertRepoData(db *sql.DB, r repo.Repo) error {
	if db == nil {
		return errors.New("nil database given")
	}

	repoID, err := getRepoID(db, r)
	if err != nil {
		return err
	}
	if repoID == nil {
		return errors.New("cannot find corresponding repository in database")
	}

	tx, err := db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	commitStmt, err := tx.Prepare(genInsQuery("commits", commitFields...) + " RETURNING id")
	if err != nil {
		return err
	}

	deltaStmt, err := tx.Prepare(genInsQuery("commit_diff_deltas", diffDeltaFields...))
	if err != nil {
		return err
	}

	for _, c := range r.GetCommits() {
		if err := insertCommit(*repoID, c, tx, commitStmt, deltaStmt); err != nil {
			return err
		}
	}

	if err := commitStmt.Close(); err != nil {
		return err
	}
	if err := deltaStmt.Close(); err != nil {
		return err
	}
	if err := tx.Commit(); err != nil {
		return err
	}

	return nil
}

// insertCommit inserts a commit into the database
func insertCommit(repoID int64, c model.Commit, tx *sql.Tx, commitStmt, deltaStmt *sql.Stmt) error {
	// TODO load user ids/email in a map once at start
	authorID, err := getUserID(tx, c.Author.Email)
	if err != nil {
		return err
	}
	committerID, err := getUserID(tx, c.Committer.Email)
	if err != nil {
		return err
	}
	hash := genCommitHash(c)

	var commitID int64
	err = commitStmt.QueryRow(
		repoID, authorID, committerID, hash,
		c.VCSID, c.Message, c.AuthorDate, c.CommitDate,
		c.FileChangedCount, c.InsertionsCount, c.DeletionsCount).Scan(&commitID)
	if err != nil {
		return err
	}

	for _, d := range c.DiffDelta {
		if err := insertDiffDelta(commitID, d, deltaStmt); err != nil {
			return err
		}
	}

	return nil
}

// insertDiffDelta inserts a commit diff delta into the database.
func insertDiffDelta(commitID int64, d model.DiffDelta, stmt *sql.Stmt) error {
	_, err := stmt.Exec(commitID, d.Status, d.Binary, d.Similarity, d.OldFilePath, d.NewFilePath)
	if err != nil {
		return err
	}
	return nil
}

// genCommitHash generates a hash (mmh3) from commit fields.
// This hash can then be used to uniquely identify a commit.
// Typically, we want to make sure not to insert twice the same commit into the
// database after an eventual second repotool run on the same repository.
func genCommitHash(c model.Commit) string {
	h := mmh3.New128()

	io.WriteString(h, c.VCSID)
	io.WriteString(h, c.Message)
	io.WriteString(h, c.Author.Name)
	io.WriteString(h, c.Author.Email)
	io.WriteString(h, c.Committer.Name)
	io.WriteString(h, c.Committer.Email)
	io.WriteString(h, c.AuthorDate.String())
	io.WriteString(h, c.CommitDate.String())
	io.WriteString(h, strconv.FormatInt(int64(c.FileChangedCount), 10))
	io.WriteString(h, strconv.FormatInt(int64(c.InsertionsCount), 10))
	io.WriteString(h, strconv.FormatInt(int64(c.DeletionsCount), 10))

	return hex.EncodeToString(h.Sum(nil))
}

// getRepoID returns the repository id of a repo in repositories table.
// If repo is not in the table, then 0 is returned.
func getRepoID(db *sql.DB, r repo.Repo) (*int64, error) {
	if db == nil {
		return nil, errors.New("nil database given")
	}

	var id *int64
	// Clone URL is unique
	err := db.QueryRow("SELECT id FROM repositories WHERE clone_url=$1", r.GetCloneURL()).Scan(&id)
	switch {
	case err == sql.ErrNoRows:
		return nil, nil
	case err != nil:
		return nil, err
	}
	return id, nil
}

// getUserID attempts to find a user ID given its email address.
// Email addresses are unique, however they may not be provided.
// If no user ID is found, nil is returned, otherwhise the user ID
// is returned.
func getUserID(tx *sql.Tx, email string) (*int64, error) {
	if tx == nil {
		return nil, errors.New("nil database given")
	}

	if len(email) == 0 {
		return nil, nil
	}

	var id *int64
	err := tx.QueryRow("SELECT id FROM users WHERE email=$1", email).Scan(&id)
	switch {
	case err == sql.ErrNoRows:
		return nil, nil
	case err != nil:
		return nil, err
	}
	return id, nil
}

// genInsQuery generates a query string for an insertion in the database.
func genInsQuery(tableName string, fields ...string) string {
	var buf bytes.Buffer

	buf.WriteString(fmt.Sprintf("INSERT INTO %s(%s)\n",
		tableName, strings.Join(fields, ",")))
	buf.WriteString("VALUES(")

	for ind := range fields {
		if ind > 0 {
			buf.WriteString(",")
		}

		buf.WriteString(fmt.Sprintf("$%d", ind+1))
	}

	buf.WriteString(")\n")

	return buf.String()
}
<commit_msg>db: store users in hashmap to avoid db queries<commit_after>// Copyright 2014-2015 The DevMine authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package repotool is able to fetch information from a source code repository.
// Typically, it can get all commits, their authors and commiters and so on
// and return this information in a JSON object. Alternatively, it is able
// to populate the information into a PostgreSQL database.
// Currently, on the Git VCS is supported.
package main

import (
	"bufio"
	"bytes"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	_ "github.com/lib/pq"
	mmh3 "github.com/spaolacci/murmur3"

	"github.com/DevMine/srcanlzr/src"

	"github.com/DevMine/repotool/config"
	"github.com/DevMine/repotool/model"
	"github.com/DevMine/repotool/repo"
)

const version = "0.1.0"

// database fields per tables
var (
	diffDeltaFields = []string{
		"commit_id",
		"file_status",
		"is_file_binary",
		"similarity",
		"old_file_path",
		"new_file_path"}

	commitFields = []string{
		"repository_id",
		"author_id",
		"committer_id",
		"hash",
		"vcs_id",
		"message",
		"author_date",
		"commit_date",
		"file_changed_count",
		"insertions_count",
		"deletions_count"}
)

func main() {
	flag.Usage = func() {
		fmt.Printf("usage: %s [OPTION(S)] [REPOSITORY PATH]\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(0)
	}

	configPath := flag.String("c", "", "configuration file")
	vflag := flag.Bool("v", false, "print version.")
	jsonflag := flag.Bool("json", true, "json output")
	dbflag := flag.Bool("db", false, "import data into the database")
	srctoolflag := flag.String("srctool", "", "read json file produced by srctool (give stdin to read from stdin)")
	flag.Parse()

	if *vflag {
		fmt.Printf("%s - %s\n", filepath.Base(os.Args[0]), version)
		os.Exit(0)
	}

	if len(flag.Args()) != 1 {
		fmt.Fprintln(os.Stderr, "invalid # of arguments")
		flag.Usage()
	}

	if *dbflag && len(*configPath) == 0 {
		fatal(errors.New("a configuration file must be specified when using db option"))
	}

	if !*jsonflag && (len(*srctoolflag) > 0) {
		fatal(errors.New("srctool flag may be used only in conjonction with json flag"))
	}

	cfg, err := config.ReadConfig(*configPath)
	if err != nil {
		fatal(err)
	}

	repoPath := flag.Arg(0)
	repository, err := repo.New(*cfg, repoPath)
	if err != nil {
		fatal(err)
	}
	defer repository.Cleanup()

	fmt.Fprintln(os.Stderr, "fetching repository commits...")
	tic := time.Now()
	err = repository.FetchCommits()
	if err != nil {
		fatal(err)
	}
	toc := time.Now()
	fmt.Fprintln(os.Stderr, "done in ", toc.Sub(tic))

	if *jsonflag && (len(*srctoolflag) == 0) {
		bs, err := json.Marshal(repository)
		if err != nil {
			fatal(err)
		}
		fmt.Println(string(bs))
	}

	if *jsonflag && (len(*srctoolflag)) > 0 {
		var r *bufio.Reader
		if *srctoolflag == strings.ToLower("stdin") {
			// read from stdin
			r = bufio.NewReader(os.Stdin)
		} else {
			// read from srctool json file
			var f *os.File
			if f, err = os.Open(*srctoolflag); err != nil {
				fatal(err)
			}
			r = bufio.NewReader(f)
		}

		buf := new(bytes.Buffer)
		if _, err = io.Copy(buf, r); err != nil {
			fatal(err)
		}

		bs := buf.Bytes()
		p, err := src.Unmarshal(bs)
		if err != nil {
			fatal(err)
		}

		p.Repo = repository.GetRepository()
		bs, err = src.Marshal(p)
		if err != nil {
			fatal(err)
		}

		fmt.Println(string(bs))
	}

	if *dbflag {
		db, err := openDBSession(cfg.Database)
		if err != nil {
			fatal(err)
		}
		defer db.Close()

		fmt.Fprintf(os.Stderr,
			"inserting %d commits from %s repository into the database...\n",
			len(repository.GetCommits()), repository.GetName())
		tic := time.Now()
		err = insertRepoData(db, repository)
		toc := time.Now()
		if err != nil {
			fmt.Fprintln(os.Stderr, err.Error())
			os.Exit(1)
		} else {
			fmt.Fprintln(os.Stderr, "done in ", toc.Sub(tic))
		}
	}
}

// fatal prints an error on standard error stream and exits.
func fatal(a ...interface{}) {
	fmt.Fprintln(os.Stderr, a...)
	os.Exit(1)
}

// fail prints an error on standard error stream.
func fail(a ...interface{}) {
	fmt.Fprintln(os.Stderr, a...)
}

// openDBSession creates a session to the database.
func openDBSession(cfg config.DatabaseConfig) (*sql.DB, error) {
	dbURL := fmt.Sprintf(
		"user='%s' password='%s' host='%s' port=%d dbname='%s' sslmode='%s'",
		cfg.UserName, cfg.Password, cfg.HostName, cfg.Port, cfg.DBName, cfg.SSLMode)

	return sql.Open("postgres", dbURL)
}

// insertRepoData inserts repository data into the database, or updates it
// if it is already there.
func insertRepoData(db *sql.DB, r repo.Repo) error {
	if db == nil {
		return errors.New("nil database given")
	}

	repoID, err := getRepoID(db, r)
	if err != nil {
		return err
	}
	if repoID == nil {
		return errors.New("cannot find corresponding repository in database")
	}

	userIDs, err := getAllUsers(db)
	if err != nil {
		return err
	}

	tx, err := db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	commitStmt, err := tx.Prepare(genInsQuery("commits", commitFields...) + " RETURNING id")
	if err != nil {
		return err
	}

	deltaStmt, err := tx.Prepare(genInsQuery("commit_diff_deltas", diffDeltaFields...))
	if err != nil {
		return err
	}

	for _, c := range r.GetCommits() {
		if err := insertCommit(userIDs, *repoID, c, tx, commitStmt, deltaStmt); err != nil {
			return err
		}
	}

	if err := commitStmt.Close(); err != nil {
		return err
	}
	if err := deltaStmt.Close(); err != nil {
		return err
	}
	if err := tx.Commit(); err != nil {
		return err
	}

	return nil
}

// insertCommit inserts a commit into the database
func insertCommit(userIDs map[string]uint64, repoID uint64, c model.Commit, tx *sql.Tx, commitStmt, deltaStmt *sql.Stmt) error {
	authorID := userIDs[c.Author.Email]
	committerID := userIDs[c.Committer.Email]
	hash := genCommitHash(c)

	var commitID uint64
	err := commitStmt.QueryRow(
		repoID, authorID, committerID, hash,
		c.VCSID, c.Message, c.AuthorDate, c.CommitDate,
		c.FileChangedCount, c.InsertionsCount, c.DeletionsCount).Scan(&commitID)
	if err != nil {
		return err
	}

	for _, d := range c.DiffDelta {
		if err := insertDiffDelta(commitID, d, deltaStmt); err != nil {
			return err
		}
	}

	return nil
}

// insertDiffDelta inserts a commit diff delta into the database.
func insertDiffDelta(commitID uint64, d model.DiffDelta, stmt *sql.Stmt) error {
	_, err := stmt.Exec(commitID, d.Status, d.Binary, d.Similarity, d.OldFilePath, d.NewFilePath)
	if err != nil {
		return err
	}
	return nil
}

// getAllUsers returns a map of all users IDs with their email address as keys.
func getAllUsers(db *sql.DB) (map[string]uint64, error) {
	rows, err := db.Query("SELECT id, email FROM users WHERE email IS NOT NULL AND email != ''")
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	userIDs := map[string]uint64{}
	for rows.Next() {
		var email string
		var id uint64
		if err := rows.Scan(&id, &email); err != nil {
			return nil, err
		}
		userIDs[email] = id
	}

	return userIDs, nil
}

// genCommitHash generates a hash (mmh3) from commit fields.
// This hash can then be used to uniquely identify a commit.
// Typically, we want to make sure not to insert twice the same commit into the
// database after an eventual second repotool run on the same repository.
func genCommitHash(c model.Commit) string {
	h := mmh3.New128()

	io.WriteString(h, c.VCSID)
	io.WriteString(h, c.Message)
	io.WriteString(h, c.Author.Name)
	io.WriteString(h, c.Author.Email)
	io.WriteString(h, c.Committer.Name)
	io.WriteString(h, c.Committer.Email)
	io.WriteString(h, c.AuthorDate.String())
	io.WriteString(h, c.CommitDate.String())
	io.WriteString(h, strconv.FormatInt(int64(c.FileChangedCount), 10))
	io.WriteString(h, strconv.FormatInt(int64(c.InsertionsCount), 10))
	io.WriteString(h, strconv.FormatInt(int64(c.DeletionsCount), 10))

	return hex.EncodeToString(h.Sum(nil))
}

// getRepoID returns the repository id of a repo in repositories table.
// If repo is not in the table, then 0 is returned.
func getRepoID(db *sql.DB, r repo.Repo) (*uint64, error) {
	if db == nil {
		return nil, errors.New("nil database given")
	}

	var id *uint64
	// Clone URL is unique
	err := db.QueryRow("SELECT id FROM repositories WHERE clone_url=$1", r.GetCloneURL()).Scan(&id)
	switch {
	case err == sql.ErrNoRows:
		return nil, nil
	case err != nil:
		return nil, err
	}
	return id, nil
}

// genInsQuery generates a query string for an insertion in the database.
func genInsQuery(tableName string, fields ...string) string {
	var buf bytes.Buffer

	buf.WriteString(fmt.Sprintf("INSERT INTO %s(%s)\n",
		tableName, strings.Join(fields, ",")))
	buf.WriteString("VALUES(")

	for ind := range fields {
		if ind > 0 {
			buf.WriteString(",")
		}

		buf.WriteString(fmt.Sprintf("$%d", ind+1))
	}

	buf.WriteString(")\n")

	return buf.String()
}
<|endoftext|>


<commit_before>package appdash

import (
	"errors"
	"fmt"
	"net/url"
	"reflect"
	"time"

	influxDBClient "github.com/influxdb/influxdb/client"
	influxDBServer "github.com/influxdb/influxdb/cmd/influxd/run"
	influxDBModels "github.com/influxdb/influxdb/models"
)

const (
	dbName               string = "appdash" // InfluxDB db name.
	spanMeasurementName  string = "spans"   // InfluxDB container name for trace spans.
	defaultTracesPerPage int    = 10        // Default number of traces per page.
)

// Compile-time "implements" check.
var _ interface {
	Store
	Queryer
} = (*InfluxDBStore)(nil)

// zeroID is ID's zero value as string.
var zeroID string = ID(0).String()

// pointFields -> influxDBClient.Point.Fields
type pointFields map[string]interface{}

type InfluxDBStore struct {
	con           *influxDBClient.Client // InfluxDB client connection.
	server        *influxDBServer.Server // InfluxDB API server.
	tracesPerPage int                    // Number of traces per page.
}

func (in *InfluxDBStore) Collect(id SpanID, anns ...Annotation) error {
	p, err := in.findSpanPoint(id)
	if err != nil {
		return err
	}

	// trace_id, span_id & parent_id are set as tags
	// because InfluxDB tags are indexed & those values
	// are used later on queries.
	tags := map[string]string{
		"trace_id":  id.Trace.String(),
		"span_id":   id.Span.String(),
		"parent_id": id.Parent.String(),
	}

	// Saving annotations as InfluxDB measurement spans fields
	// which are not indexed.
	fields := make(map[string]interface{}, len(anns))
	for _, ann := range anns {
		fields[ann.Key] = string(ann.Value)
	}

	if p != nil { // span exists on DB.
		p.Measurement = spanMeasurementName
		p.Tags = tags
		// Extends fields in order to preserve the p.Fields
		// already saved on DB.
		p.Fields = extendFields(fields, p.Fields)
	} else { // new span to be saved on DB.
		p = &influxDBClient.Point{
			Measurement: spanMeasurementName,
			Tags:        tags,   // indexed metadata.
			Fields:      fields, // non-indexed metadata.
			Time:        time.Now().UTC(),
		}
	}

	// InfluxDB point represents a single span.
	pts := []influxDBClient.Point{*p}
	bps := influxDBClient.BatchPoints{
		Points:          pts,
		Database:        dbName,
		RetentionPolicy: "default",
	}
	_, writeErr := in.con.Write(bps)
	if writeErr != nil {
		return writeErr
	}
	return nil
}

func (in *InfluxDBStore) Trace(id ID) (*Trace, error) {
	trace := &Trace{}

	// GROUP BY * -> meaning group by all tags(trace_id, span_id & parent_id)
	// grouping by all tags includes those and it's values on the query response.
	q := fmt.Sprintf("SELECT * FROM spans WHERE trace_id='%s' GROUP BY *", id)
	result, err := in.executeOneQuery(q)
	if err != nil {
		return nil, err
	}

	// result.Series -> A slice containing all the spans.
	if len(result.Series) == 0 {
		return nil, errors.New("trace not found")
	}

	// Iterate over series(spans) to create trace children's & set trace fields.
	var rootSpanSet bool
	for _, s := range result.Series {
		var isRootSpan bool
		span, err := newSpanFromRow(&s)
		if err != nil {
			return nil, err
		}
		annotations, err := annotationsFromRow(&s)
		if err != nil {
			return trace, nil
		}
		span.Annotations = *annotations
		if span.ID.IsRoot() && rootSpanSet {
			return nil, errors.New("unexpected multiple root spans")
		}
		if span.ID.IsRoot() && !rootSpanSet {
			isRootSpan = true
		}
		if isRootSpan { // root span.
			trace.Span = *span
			rootSpanSet = true
		} else { // children span.
			trace.Sub = append(trace.Sub, &Trace{Span: *span})
		}
	}
	return trace, nil
}

func (in *InfluxDBStore) Traces() ([]*Trace, error) {
	traces := make([]*Trace, 0)

	// GROUP BY * -> meaning group by all tags(trace_id, span_id & parent_id)
	// grouping by all tags includes those and it's values on the query response.
	rootSpansQuery := fmt.Sprintf("SELECT * FROM spans WHERE parent_id='%s' GROUP BY * LIMIT %d", zeroID, in.tracesPerPage)
	rootSpansResult, err := in.executeOneQuery(rootSpansQuery)
	if err != nil {
		return nil, err
	}

	// result.Series -> A slice containing all the spans.
	if len(rootSpansResult.Series) == 0 {
		return traces, nil
	}

	// Cache to keep track of traces to be returned.
	tracesCache := make(map[ID]*Trace, 0)

	// Iterate over series(spans) to create traces.
	for _, s := range rootSpansResult.Series {
		span, err := newSpanFromRow(&s)
		if err != nil {
			return nil, err
		}
		annotations, err := annotationsFromRow(&s)
		if err != nil {
			return nil, err
		}
		span.Annotations = *annotations
		_, present := tracesCache[span.ID.Trace]
		if !present {
			tracesCache[span.ID.Trace] = &Trace{Span: *span}
		} else {
			return nil, errors.New("duplicated root span")
		}
	}

	// Using 'OR' since 'IN' not supported yet.
	where := `WHERE `
	var i int = 1
	for _, trace := range tracesCache {
		where += fmt.Sprintf("(trace_id='%s' AND parent_id!='%s')", trace.Span.ID.Trace, zeroID)

		// Adds 'OR' except for last iteration.
		if i != len(tracesCache) && len(tracesCache) > 1 {
			where += " OR "
		}
		i += 1
	}

	// Queries for all children spans of the traces to be returned.
	childrenSpansQuery := fmt.Sprintf("SELECT * FROM spans %s GROUP BY *", where)
	childrenSpansResult, err := in.executeOneQuery(childrenSpansQuery)
	if err != nil {
		return nil, err
	}

	// Iterate over series(children spans) to create sub-traces
	// and associates sub-traces with it's parent trace.
	for _, s := range childrenSpansResult.Series {
		span, err := newSpanFromRow(&s)
		if err != nil {
			return nil, err
		}
		annotations, err := annotationsFromRow(&s)
		if err != nil {
			return nil, err
		}
		span.Annotations = *annotations
		trace, present := tracesCache[span.ID.Trace]
		if !present { // Root trace not added.
			return nil, errors.New("parent not found")
		} else { // Root trace already added so append a sub-trace.
			trace.Sub = append(trace.Sub, &Trace{Span: *span})
		}
	}
	for _, trace := range tracesCache {
		traces = append(traces, trace)
	}
	return traces, nil
}

func (in *InfluxDBStore) Close() {
	in.server.Close()
}

func (in *InfluxDBStore) createDBIfNotExists() error {
	// If no errors query execution was successfully - either DB was created or already exists.
	response, err := in.con.Query(influxDBClient.Query{
		Command: fmt.Sprintf("%s %s", "CREATE DATABASE IF NOT EXISTS", dbName),
	})
	if err != nil {
		return err
	}
	if response.Error() != nil {
		return response.Error()
	}
	return nil
}

func (in *InfluxDBStore) executeOneQuery(command string) (*influxDBClient.Result, error) {
	response, err := in.con.Query(influxDBClient.Query{
		Command:  command,
		Database: dbName,
	})
	if err != nil {
		return nil, err
	}
	if response.Error() != nil {
		return nil, response.Error()
	}

	// Expecting one result, since a single query is executed.
	if len(response.Results) != 1 {
		return nil, errors.New("unexpected number of results for an influxdb single query")
	}
	return &response.Results[0], nil
}

func (in *InfluxDBStore) findSpanPoint(ID SpanID) (*influxDBClient.Point, error) {
	q := fmt.Sprintf(`
		SELECT * FROM spans WHERE trace_id='%s' AND span_id='%s' AND parent_id='%s' GROUP BY *
	`, ID.Trace, ID.Span, ID.Parent)
	result, err := in.executeOneQuery(q)
	if err != nil {
		return nil, err
	}
	if len(result.Series) == 0 {
		return nil, nil
	}
	if len(result.Series) > 1 {
		return nil, errors.New("unexpected multiple series")
	}
	r := result.Series[0]
	if len(r.Values) == 0 {
		return nil, errors.New("unexpected empty series")
	}
	p := influxDBClient.Point{
		Fields: make(pointFields, 0),
	}
	fields := r.Values[0]
	for i, field := range fields {
		key := r.Columns[i]
		switch field.(type) {
		case string:
			// time field is set by InfluxDB not related to annotations.
			if key == "time" {
				t, err := time.Parse(time.RFC3339Nano, field.(string))
				if err != nil {
					return nil, err
				}
				p.Time = t
			}
			if key == "Name" {
				p.Fields[key] = field.(string)
			}
		case nil:
			continue
		default:
			return nil, fmt.Errorf("unexpected time type: %v", reflect.TypeOf(field))
		}
	}
	return &p, err
}

func (in *InfluxDBStore) init(server *influxDBServer.Server) error {
	in.server = server
	url, err := url.Parse(fmt.Sprintf("http://%s:%d", influxDBClient.DefaultHost, influxDBClient.DefaultPort))
	if err != nil {
		return err
	}
	con, err := influxDBClient.NewClient(influxDBClient.Config{URL: *url})
	if err != nil {
		return err
	}
	in.con = con
	if err := in.createDBIfNotExists(); err != nil {
		return err
	}
	in.tracesPerPage = defaultTracesPerPage
	return nil
}

func annotationsFromRow(r *influxDBModels.Row) (*Annotations, error) {
	// Actually an influxDBModels.Row represents a single InfluxDB serie.
	// r.Values[n] is a slice containing span's annotation values.
	var fields []interface{}
	if len(r.Values) == 1 {
		fields = r.Values[0]
	}

	// len(r.Values) might be greater than one - meaning there are
	// some spans to drop, see: InfluxDBStore.Collect(...).
	// If so last one is picked.
	if len(r.Values) > 1 {
		fields = r.Values[len(r.Values)-1]
	}
	annotations := make(Annotations, 0)

	// Iterates over fields which represent span's annotation values.
	for i, field := range fields {
		// It is safe to do column[0] (eg. 'Server.Request.Method')
		// matches fields[0] (eg. 'GET')
		key := r.Columns[i]
		var value []byte
		switch field.(type) {
		case string:
			value = []byte(field.(string))
		case nil:
		default:
			return nil, fmt.Errorf("unexpected field type: %v", reflect.TypeOf(field))
		}
		a := Annotation{
			Key:   key,
			Value: value,
		}
		annotations = append(annotations, a)
	}

	return &annotations, nil
}

// extendFields replaces existing items on dst from src.
func extendFields(dst, src pointFields) pointFields {
	for k, v := range src {
		if _, present := dst[k]; present {
			dst[k] = v
		}
	}
	return dst
}

func newSpanFromRow(r *influxDBModels.Row) (*Span, error) {
	span := &Span{}
	traceID, err := ParseID(r.Tags["trace_id"])
	if err != nil {
		return nil, err
	}
	spanID, err := ParseID(r.Tags["span_id"])
	if err != nil {
		return nil, err
	}
	parentID, err := ParseID(r.Tags["parent_id"])
	if err != nil {
		return nil, err
	}
	span.ID = SpanID{
		Trace:  ID(traceID),
		Span:   ID(spanID),
		Parent: ID(parentID),
	}
	return span, nil
}

type InfluxDBStoreConfig struct {
	Server    *influxDBServer.Config
	BuildInfo *influxDBServer.BuildInfo
}

func NewInfluxDBStore(config InfluxDBStoreConfig) (*InfluxDBStore, error) {
	//TODO: add Authentication.
	s, err := influxDBServer.NewServer(config.Server, config.BuildInfo)
	if err != nil {
		return nil, err
	}
	if err := s.Open(); err != nil {
		return nil, err
	}
	var in InfluxDBStore
	if err := in.init(s); err != nil {
		return nil, err
	}
	return &in, nil
}
<commit_msg>set all other fields diff than Name too<commit_after>package appdash

import (
	"errors"
	"fmt"
	"net/url"
	"reflect"
	"time"

	influxDBClient "github.com/influxdb/influxdb/client"
	influxDBServer "github.com/influxdb/influxdb/cmd/influxd/run"
	influxDBModels "github.com/influxdb/influxdb/models"
)

const (
	dbName               string = "appdash" // InfluxDB db name.
	spanMeasurementName  string = "spans"   // InfluxDB container name for trace spans.
	defaultTracesPerPage int    = 10        // Default number of traces per page.
)

// Compile-time "implements" check.
var _ interface {
	Store
	Queryer
} = (*InfluxDBStore)(nil)

// zeroID is ID's zero value as string.
var zeroID string = ID(0).String()

// pointFields -> influxDBClient.Point.Fields
type pointFields map[string]interface{}

type InfluxDBStore struct {
	con           *influxDBClient.Client // InfluxDB client connection.
	server        *influxDBServer.Server // InfluxDB API server.
	tracesPerPage int                    // Number of traces per page.
}

func (in *InfluxDBStore) Collect(id SpanID, anns ...Annotation) error {
	p, err := in.findSpanPoint(id)
	if err != nil {
		return err
	}

	// trace_id, span_id & parent_id are set as tags
	// because InfluxDB tags are indexed & those values
	// are used later on queries.
	tags := map[string]string{
		"trace_id":  id.Trace.String(),
		"span_id":   id.Span.String(),
		"parent_id": id.Parent.String(),
	}

	// Saving annotations as InfluxDB measurement spans fields
	// which are not indexed.
	fields := make(map[string]interface{}, len(anns))
	for _, ann := range anns {
		fields[ann.Key] = string(ann.Value)
	}

	if p != nil { // span exists on DB.
		p.Measurement = spanMeasurementName
		p.Tags = tags
		// Extends fields in order to preserve the p.Fields
		// already saved on DB.
		p.Fields = extendFields(fields, p.Fields)
	} else { // new span to be saved on DB.
		p = &influxDBClient.Point{
			Measurement: spanMeasurementName,
			Tags:        tags,   // indexed metadata.
			Fields:      fields, // non-indexed metadata.
			Time:        time.Now().UTC(),
		}
	}

	// InfluxDB point represents a single span.
	pts := []influxDBClient.Point{*p}
	bps := influxDBClient.BatchPoints{
		Points:          pts,
		Database:        dbName,
		RetentionPolicy: "default",
	}
	_, writeErr := in.con.Write(bps)
	if writeErr != nil {
		return writeErr
	}
	return nil
}

func (in *InfluxDBStore) Trace(id ID) (*Trace, error) {
	trace := &Trace{}

	// GROUP BY * -> meaning group by all tags(trace_id, span_id & parent_id)
	// grouping by all tags includes those and it's values on the query response.
	q := fmt.Sprintf("SELECT * FROM spans WHERE trace_id='%s' GROUP BY *", id)
	result, err := in.executeOneQuery(q)
	if err != nil {
		return nil, err
	}

	// result.Series -> A slice containing all the spans.
	if len(result.Series) == 0 {
		return nil, errors.New("trace not found")
	}

	// Iterate over series(spans) to create trace children's & set trace fields.
	var rootSpanSet bool
	for _, s := range result.Series {
		var isRootSpan bool
		span, err := newSpanFromRow(&s)
		if err != nil {
			return nil, err
		}
		annotations, err := annotationsFromRow(&s)
		if err != nil {
			return trace, nil
		}
		span.Annotations = *annotations
		if span.ID.IsRoot() && rootSpanSet {
			return nil, errors.New("unexpected multiple root spans")
		}
		if span.ID.IsRoot() && !rootSpanSet {
			isRootSpan = true
		}
		if isRootSpan { // root span.
			trace.Span = *span
			rootSpanSet = true
		} else { // children span.
			trace.Sub = append(trace.Sub, &Trace{Span: *span})
		}
	}
	return trace, nil
}

func (in *InfluxDBStore) Traces() ([]*Trace, error) {
	traces := make([]*Trace, 0)

	// GROUP BY * -> meaning group by all tags(trace_id, span_id & parent_id)
	// grouping by all tags includes those and it's values on the query response.
	rootSpansQuery := fmt.Sprintf("SELECT * FROM spans WHERE parent_id='%s' GROUP BY * LIMIT %d", zeroID, in.tracesPerPage)
	rootSpansResult, err := in.executeOneQuery(rootSpansQuery)
	if err != nil {
		return nil, err
	}

	// result.Series -> A slice containing all the spans.
	if len(rootSpansResult.Series) == 0 {
		return traces, nil
	}

	// Cache to keep track of traces to be returned.
	tracesCache := make(map[ID]*Trace, 0)

	// Iterate over series(spans) to create traces.
	for _, s := range rootSpansResult.Series {
		span, err := newSpanFromRow(&s)
		if err != nil {
			return nil, err
		}
		annotations, err := annotationsFromRow(&s)
		if err != nil {
			return nil, err
		}
		span.Annotations = *annotations
		_, present := tracesCache[span.ID.Trace]
		if !present {
			tracesCache[span.ID.Trace] = &Trace{Span: *span}
		} else {
			return nil, errors.New("duplicated root span")
		}
	}

	// Using 'OR' since 'IN' not supported yet.
	where := `WHERE `
	var i int = 1
	for _, trace := range tracesCache {
		where += fmt.Sprintf("(trace_id='%s' AND parent_id!='%s')", trace.Span.ID.Trace, zeroID)

		// Adds 'OR' except for last iteration.
		if i != len(tracesCache) && len(tracesCache) > 1 {
			where += " OR "
		}
		i += 1
	}

	// Queries for all children spans of the traces to be returned.
	childrenSpansQuery := fmt.Sprintf("SELECT * FROM spans %s GROUP BY *", where)
	childrenSpansResult, err := in.executeOneQuery(childrenSpansQuery)
	if err != nil {
		return nil, err
	}

	// Iterate over series(children spans) to create sub-traces
	// and associates sub-traces with it's parent trace.
	for _, s := range childrenSpansResult.Series {
		span, err := newSpanFromRow(&s)
		if err != nil {
			return nil, err
		}
		annotations, err := annotationsFromRow(&s)
		if err != nil {
			return nil, err
		}
		span.Annotations = *annotations
		trace, present := tracesCache[span.ID.Trace]
		if !present { // Root trace not added.
			return nil, errors.New("parent not found")
		} else { // Root trace already added so append a sub-trace.
			trace.Sub = append(trace.Sub, &Trace{Span: *span})
		}
	}
	for _, trace := range tracesCache {
		traces = append(traces, trace)
	}
	return traces, nil
}

func (in *InfluxDBStore) Close() {
	in.server.Close()
}

func (in *InfluxDBStore) createDBIfNotExists() error {
	// If no errors query execution was successfully - either DB was created or already exists.
	response, err := in.con.Query(influxDBClient.Query{
		Command: fmt.Sprintf("%s %s", "CREATE DATABASE IF NOT EXISTS", dbName),
	})
	if err != nil {
		return err
	}
	if response.Error() != nil {
		return response.Error()
	}
	return nil
}

func (in *InfluxDBStore) executeOneQuery(command string) (*influxDBClient.Result, error) {
	response, err := in.con.Query(influxDBClient.Query{
		Command:  command,
		Database: dbName,
	})
	if err != nil {
		return nil, err
	}
	if response.Error() != nil {
		return nil, response.Error()
	}

	// Expecting one result, since a single query is executed.
	if len(response.Results) != 1 {
		return nil, errors.New("unexpected number of results for an influxdb single query")
	}
	return &response.Results[0], nil
}

func (in *InfluxDBStore) findSpanPoint(ID SpanID) (*influxDBClient.Point, error) {
	q := fmt.Sprintf(`
		SELECT * FROM spans WHERE trace_id='%s' AND span_id='%s' AND parent_id='%s' GROUP BY *
	`, ID.Trace, ID.Span, ID.Parent)
	result, err := in.executeOneQuery(q)
	if err != nil {
		return nil, err
	}
	if len(result.Series) == 0 {
		return nil, nil
	}
	if len(result.Series) > 1 {
		return nil, errors.New("unexpected multiple series")
	}
	r := result.Series[0]
	if len(r.Values) == 0 {
		return nil, errors.New("unexpected empty series")
	}
	p := influxDBClient.Point{
		Fields: make(pointFields, 0),
	}
	fields := r.Values[0]
	for i, field := range fields {
		key := r.Columns[i]
		switch field.(type) {
		case string:
			// time field is set by InfluxDB not related to annotations.
			if key == "time" {
				t, err := time.Parse(time.RFC3339Nano, field.(string))
				if err != nil {
					return nil, err
				}
				p.Time = t
			}
			p.Fields[key] = field.(string)
		case nil:
			continue
		default:
			return nil, fmt.Errorf("unexpected field type: %v", reflect.TypeOf(field))
		}
	}
	return &p, err
}

func (in *InfluxDBStore) init(server *influxDBServer.Server) error {
	in.server = server
	url, err := url.Parse(fmt.Sprintf("http://%s:%d", influxDBClient.DefaultHost, influxDBClient.DefaultPort))
	if err != nil {
		return err
	}
	con, err := influxDBClient.NewClient(influxDBClient.Config{URL: *url})
	if err != nil {
		return err
	}
	in.con = con
	if err := in.createDBIfNotExists(); err != nil {
		return err
	}
	in.tracesPerPage = defaultTracesPerPage
	return nil
}

func annotationsFromRow(r *influxDBModels.Row) (*Annotations, error) {
	// Actually an influxDBModels.Row represents a single InfluxDB serie.
	// r.Values[n] is a slice containing span's annotation values.
	var fields []interface{}
	if len(r.Values) == 1 {
		fields = r.Values[0]
	}

	// len(r.Values) might be greater than one - meaning there are
	// some spans to drop, see: InfluxDBStore.Collect(...).
	// If so last one is picked.
	if len(r.Values) > 1 {
		fields = r.Values[len(r.Values)-1]
	}
	annotations := make(Annotations, 0)

	// Iterates over fields which represent span's annotation values.
	for i, field := range fields {
		// It is safe to do column[0] (eg. 'Server.Request.Method')
		// matches fields[0] (eg. 'GET')
		key := r.Columns[i]
		var value []byte
		switch field.(type) {
		case string:
			value = []byte(field.(string))
		case nil:
		default:
			return nil, fmt.Errorf("unexpected field type: %v", reflect.TypeOf(field))
		}
		a := Annotation{
			Key:   key,
			Value: value,
		}
		annotations = append(annotations, a)
	}

	return &annotations, nil
}

// extendFields replaces existing items on dst from src.
func extendFields(dst, src pointFields) pointFields {
	for k, v := range src {
		if _, present := dst[k]; present {
			dst[k] = v
		}
	}
	return dst
}

func newSpanFromRow(r *influxDBModels.Row) (*Span, error) {
	span := &Span{}
	traceID, err := ParseID(r.Tags["trace_id"])
	if err != nil {
		return nil, err
	}
	spanID, err := ParseID(r.Tags["span_id"])
	if err != nil {
		return nil, err
	}
	parentID, err := ParseID(r.Tags["parent_id"])
	if err != nil {
		return nil, err
	}
	span.ID = SpanID{
		Trace:  ID(traceID),
		Span:   ID(spanID),
		Parent: ID(parentID),
	}
	return span, nil
}

type InfluxDBStoreConfig struct {
	Server    *influxDBServer.Config
	BuildInfo *influxDBServer.BuildInfo
}

func NewInfluxDBStore(config InfluxDBStoreConfig) (*InfluxDBStore, error) {
	//TODO: add Authentication.
	s, err := influxDBServer.NewServer(config.Server, config.BuildInfo)
	if err != nil {
		return nil, err
	}
	if err := s.Open(); err != nil {
		return nil, err
	}
	var in InfluxDBStore
	if err := in.init(s); err != nil {
		return nil, err
	}
	return &in, nil
}
<|endoftext|>


<commit_before>package saml2

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/tls"
	"encoding/xml"
	"fmt"
	"io"
)

//Response is an abstraction type for handling the information in a SAML
//assertion
type Response struct {
	Destination string       `xml:"Destination,attr"`
	Issuer      string       `xml:"Issuer"`
	Value       string       `xml:",attr"`
	Key         EncryptedKey `xml:"EncryptedAssertion>EncryptedData>KeyInfo>EncryptedKey"`
	Data        string       `xml:"EncryptedAssertion>EncryptedData>CipherData>CipherValue"`
	Signature   string       `xml:"Signature>SignatureValue"`
	Digest      string       `xml:"Signature>SignedInfo>Reference>DigestValue"`
}

//NewResponseFromReader returns a Response or error based on the given reader.
func NewResponseFromReader(r io.Reader) (*Response, error) {
	buf := &bytes.Buffer{}

	var res Response

	//Decode and copy bytes into buffer
	err := xml.NewDecoder(io.TeeReader(r, buf)).Decode(&res)
	if err != nil {
		return nil, err
	}

	return &res, nil
}

//Decrypt returns the byte slice contained in the encrypted data.
func (sr *Response) Decrypt(cert tls.Certificate) ([]byte, error) {
	data, err := xmlBytes(sr.Data)
	if err != nil {
		return nil, err
	}

	k, err := sr.Key.DecryptSymmetricKey(cert)

	if err != nil {
		return nil, fmt.Errorf("cannot decrypt, error retrieving private key: %v", err)
	}

	plainText := make([]byte, len(data))

	//Get CBC decrypter using IV
	c := cipher.NewCBCDecrypter(k, data[:aes.BlockSize])

	//Decrypt blocks
	c.CryptBlocks(plainText, data[aes.BlockSize:])

	//Remove zero block if needed
	plainText = bytes.TrimRight(plainText, string([]byte{0}))

	//Calculate index tot remove based on padding
	padLength := plainText[len(plainText)-1]
	lastGoodIndex := len(plainText) - int(padLength)

	return plainText[:lastGoodIndex], nil
}
<commit_msg>Add support for AES-GCM.<commit_after>package saml2

import (
	"bytes"
	"crypto/cipher"
	"crypto/tls"
	"encoding/xml"
	"fmt"
	"io"
)

//Response is an abstraction type for handling the information in a SAML
//assertion
type Response struct {
	Destination      string           `xml:"Destination,attr"`
	Issuer           string           `xml:"Issuer"`
	Value            string           `xml:",attr"`
	EncryptionMethod EncryptionMethod `xml:"EncryptedAssertion>EncryptedData>EncryptionMethod"`
	Key              EncryptedKey     `xml:"EncryptedAssertion>EncryptedData>KeyInfo>EncryptedKey"`
	Data             string           `xml:"EncryptedAssertion>EncryptedData>CipherData>CipherValue"`
	Signature        string           `xml:"Signature>SignatureValue"`
	Digest           string           `xml:"Signature>SignedInfo>Reference>DigestValue"`
}

//NewResponseFromReader returns a Response or error based on the given reader.
func NewResponseFromReader(r io.Reader) (*Response, error) {
	buf := &bytes.Buffer{}

	var res Response

	//Decode and copy bytes into buffer
	err := xml.NewDecoder(io.TeeReader(r, buf)).Decode(&res)
	if err != nil {
		return nil, err
	}

	return &res, nil
}

//Decrypt returns the byte slice contained in the encrypted data.
func (sr *Response) Decrypt(cert tls.Certificate) ([]byte, error) {
	data, err := xmlBytes(sr.Data)
	if err != nil {
		return nil, err
	}

	k, err := sr.Key.DecryptSymmetricKey(cert)
	if err != nil {
		return nil, fmt.Errorf("cannot decrypt, error retrieving private key: %s", err)
	}

	switch sr.EncryptionMethod.Algorithm {
	case MethodAES128GCM:
		c, err := cipher.NewGCM(k)
		if err != nil {
			return nil, fmt.Errorf("cannot create AES-GCM: %s", err)
		}

		nonce, data := data[:c.NonceSize()], data[c.NonceSize():]
		plainText, err := c.Open(nil, nonce, data, nil)
		if err != nil {
			return nil, fmt.Errorf("cannot open AES-GCM: %s", err)
		}
		return plainText, nil
	case MethodAES128CBC:
		nonce, data := data[:k.BlockSize()], data[k.BlockSize():]
		c := cipher.NewCBCDecrypter(k, nonce)
		c.CryptBlocks(data, data)

		// Remove zero bytes
		data = bytes.TrimRight(data, "\x00")

		// Calculate index to remove based on padding
		padLength := data[len(data)-1]
		lastGoodIndex := len(data) - int(padLength)
		return data[:lastGoodIndex], nil
	default:
		return nil, fmt.Errorf("unknown symmetric encryption method %#v", sr.EncryptionMethod.Algorithm)
	}
}
<|endoftext|>


<commit_before>package influxql

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"sort"

	"github.com/gogo/protobuf/proto"
	internal "github.com/influxdata/influxdb/influxql/internal"
)

// ZeroTime is the Unix nanosecond timestamp for time.Time{}.
const ZeroTime = int64(-6795364578871345152)

// Point represents a value in a series that occurred at a given time.
type Point interface {
	// Name and tags uniquely identify the series the value belongs to.
	name() string
	tags() Tags

	// The time that the value occurred at.
	time() int64

	// The value at the given time.
	value() interface{}

	// Auxillary values passed along with the value.
	aux() []interface{}
}

// Points represents a list of points.
type Points []Point

// Clone returns a deep copy of a.
func (a Points) Clone() []Point {
	other := make([]Point, len(a))
	for i, p := range a {
		if p == nil {
			other[i] = nil
			continue
		}

		switch p := p.(type) {
		case *FloatPoint:
			other[i] = p.Clone()
		case *IntegerPoint:
			other[i] = p.Clone()
		case *StringPoint:
			other[i] = p.Clone()
		case *BooleanPoint:
			other[i] = p.Clone()
		default:
			panic(fmt.Sprintf("unable to clone point: %T", p))
		}
	}
	return other
}

// Tags represent a map of keys and values.
// It memoizes its key so it can be used efficiently during query execution.
type Tags struct {
	id string
	m  map[string]string
}

// NewTags returns a new instance of Tags.
func NewTags(m map[string]string) Tags {
	if len(m) == 0 {
		return Tags{}
	}
	return Tags{
		id: string(encodeTags(m)),
		m:  m,
	}
}

// newTagsID returns a new instance of Tags parses from a tag id.
func newTagsID(id string) Tags {
	m := decodeTags([]byte(id))
	if len(m) == 0 {
		return Tags{}
	}
	return Tags{id: id, m: m}
}

// ID returns the string identifier for the tags.
func (t Tags) ID() string { return t.id }

// KeyValues returns the underlying map for the tags.
func (t Tags) KeyValues() map[string]string { return t.m }

// Keys returns a sorted list of all keys on the tag.
func (t *Tags) Keys() []string {
	if t == nil {
		return nil
	}

	var a []string
	for k := range t.m {
		a = append(a, k)
	}
	sort.Strings(a)
	return a
}

// Value returns the value for a given key.
func (t *Tags) Value(k string) string {
	if t == nil {
		return ""
	}
	return t.m[k]
}

// Subset returns a new tags object with a subset of the keys.
func (t *Tags) Subset(keys []string) Tags {
	if len(keys) == 0 {
		return Tags{}
	}

	// If keys match existing keys, simply return this tagset.
	if keysMatch(t.m, keys) {
		return *t
	}

	// Otherwise create new tag set.
	m := make(map[string]string, len(keys))
	for _, k := range keys {
		m[k] = t.m[k]
	}
	return NewTags(m)
}

// Equals returns true if t equals other.
func (t *Tags) Equals(other *Tags) bool {
	if t == nil && other == nil {
		return true
	} else if t == nil || other == nil {
		return false
	}
	return t.id == other.id
}

// keysMatch returns true if m has exactly the same keys as listed in keys.
func keysMatch(m map[string]string, keys []string) bool {
	if len(keys) != len(m) {
		return false
	}

	for _, k := range keys {
		if _, ok := m[k]; !ok {
			return false
		}
	}

	return true
}

// encodeTags converts a map of strings to an identifier.
func encodeTags(m map[string]string) []byte {
	// Empty maps marshal to empty bytes.
	if len(m) == 0 {
		return nil
	}

	// Extract keys and determine final size.
	sz := (len(m) * 2) - 1 // separators
	keys := make([]string, 0, len(m))
	for k, v := range m {
		keys = append(keys, k)
		sz += len(k) + len(v)
	}
	sort.Strings(keys)

	// Generate marshaled bytes.
	b := make([]byte, sz)
	buf := b
	for _, k := range keys {
		copy(buf, k)
		buf[len(k)] = '\x00'
		buf = buf[len(k)+1:]
	}
	for i, k := range keys {
		v := m[k]
		copy(buf, v)
		if i < len(keys)-1 {
			buf[len(v)] = '\x00'
			buf = buf[len(v)+1:]
		}
	}
	return b
}

// decodeTags parses an identifier into a map of tags.
func decodeTags(id []byte) map[string]string {
	a := bytes.Split(id, []byte{'\x00'})

	// There must be an even number of segments.
	if len(a) > 0 && len(a)%2 == 1 {
		a = a[:len(a)-1]
	}

	// Return nil if there are no segments.
	if len(a) == 0 {
		return nil
	}

	// Decode key/value tags.
	m := make(map[string]string)
	for i := 0; i < len(a); i += 2 {
		m[string(a[i])] = string(a[i+1])
	}
	return m
}

func encodeAux(aux []interface{}) []*internal.Aux {
	pb := make([]*internal.Aux, len(aux))
	for i := range aux {
		switch v := aux[i].(type) {
		case float64:
			pb[i] = &internal.Aux{DataType: proto.Int32(Float), FloatValue: proto.Float64(v)}
		case *float64:
			pb[i] = &internal.Aux{DataType: proto.Int32(Float)}
		case int64:
			pb[i] = &internal.Aux{DataType: proto.Int32(Integer), IntegerValue: proto.Int64(v)}
		case *int64:
			pb[i] = &internal.Aux{DataType: proto.Int32(Integer)}
		case string:
			pb[i] = &internal.Aux{DataType: proto.Int32(String), StringValue: proto.String(v)}
		case *string:
			pb[i] = &internal.Aux{DataType: proto.Int32(String)}
		case bool:
			pb[i] = &internal.Aux{DataType: proto.Int32(Boolean), BooleanValue: proto.Bool(v)}
		case *bool:
			pb[i] = &internal.Aux{DataType: proto.Int32(Boolean)}
		default:
			pb[i] = &internal.Aux{DataType: proto.Int32(int32(Unknown))}
		}
	}
	return pb
}

func decodeAux(pb []*internal.Aux) []interface{} {
	if len(pb) == 0 {
		return nil
	}

	aux := make([]interface{}, len(pb))
	for i := range pb {
		switch pb[i].GetDataType() {
		case Float:
			if pb[i].FloatValue != nil {
				aux[i] = *pb[i].FloatValue
			} else {
				aux[i] = (*float64)(nil)
			}
		case Integer:
			if pb[i].IntegerValue != nil {
				aux[i] = *pb[i].IntegerValue
			} else {
				aux[i] = (*int64)(nil)
			}
		case String:
			if pb[i].StringValue != nil {
				aux[i] = *pb[i].StringValue
			} else {
				aux[i] = (*string)(nil)
			}
		case Boolean:
			if pb[i].BooleanValue != nil {
				aux[i] = *pb[i].BooleanValue
			} else {
				aux[i] = (*bool)(nil)
			}
		default:
			aux[i] = nil
		}
	}
	return aux
}

// NewPointDecoder decodes generic points from a reader.
type PointDecoder struct {
	r     io.Reader
	stats IteratorStats
}

// NewPointDecoder returns a new instance of PointDecoder that reads from r.
func NewPointDecoder(r io.Reader) *PointDecoder {
	return &PointDecoder{r: r}
}

// Stats returns iterator stats embedded within the stream.
func (dec *PointDecoder) Stats() IteratorStats { return dec.stats }

// DecodePoint reads from the underlying reader and unmarshals into p.
func (dec *PointDecoder) DecodePoint(p *Point) error {
	for {
		// Read length.
		var sz uint32
		if err := binary.Read(dec.r, binary.BigEndian, &sz); err != nil {
			return err
		}

		// Read point data.
		buf := make([]byte, sz)
		if _, err := io.ReadFull(dec.r, buf); err != nil {
			return err
		}

		// Unmarshal into point.
		var pb internal.Point
		if err := proto.Unmarshal(buf, &pb); err != nil {
			return err
		}

		// If the point contains stats then read stats and retry.
		if pb.Stats != nil {
			dec.stats = decodeIteratorStats(pb.Stats)
			continue
		}

		if pb.IntegerValue != nil {
			*p = decodeIntegerPoint(&pb)
		} else if pb.StringValue != nil {
			*p = decodeStringPoint(&pb)
		} else if pb.BooleanValue != nil {
			*p = decodeBooleanPoint(&pb)
		} else {
			*p = decodeFloatPoint(&pb)
		}

		return nil
	}
}
<commit_msg>Fix decodeTags to decode tags correctly from encodeTags<commit_after>package influxql

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"sort"

	"github.com/gogo/protobuf/proto"
	internal "github.com/influxdata/influxdb/influxql/internal"
)

// ZeroTime is the Unix nanosecond timestamp for time.Time{}.
const ZeroTime = int64(-6795364578871345152)

// Point represents a value in a series that occurred at a given time.
type Point interface {
	// Name and tags uniquely identify the series the value belongs to.
	name() string
	tags() Tags

	// The time that the value occurred at.
	time() int64

	// The value at the given time.
	value() interface{}

	// Auxillary values passed along with the value.
	aux() []interface{}
}

// Points represents a list of points.
type Points []Point

// Clone returns a deep copy of a.
func (a Points) Clone() []Point {
	other := make([]Point, len(a))
	for i, p := range a {
		if p == nil {
			other[i] = nil
			continue
		}

		switch p := p.(type) {
		case *FloatPoint:
			other[i] = p.Clone()
		case *IntegerPoint:
			other[i] = p.Clone()
		case *StringPoint:
			other[i] = p.Clone()
		case *BooleanPoint:
			other[i] = p.Clone()
		default:
			panic(fmt.Sprintf("unable to clone point: %T", p))
		}
	}
	return other
}

// Tags represent a map of keys and values.
// It memoizes its key so it can be used efficiently during query execution.
type Tags struct {
	id string
	m  map[string]string
}

// NewTags returns a new instance of Tags.
func NewTags(m map[string]string) Tags {
	if len(m) == 0 {
		return Tags{}
	}
	return Tags{
		id: string(encodeTags(m)),
		m:  m,
	}
}

// newTagsID returns a new instance of Tags parses from a tag id.
func newTagsID(id string) Tags {
	m := decodeTags([]byte(id))
	if len(m) == 0 {
		return Tags{}
	}
	return Tags{id: id, m: m}
}

// ID returns the string identifier for the tags.
func (t Tags) ID() string { return t.id }

// KeyValues returns the underlying map for the tags.
func (t Tags) KeyValues() map[string]string { return t.m }

// Keys returns a sorted list of all keys on the tag.
func (t *Tags) Keys() []string {
	if t == nil {
		return nil
	}

	var a []string
	for k := range t.m {
		a = append(a, k)
	}
	sort.Strings(a)
	return a
}

// Value returns the value for a given key.
func (t *Tags) Value(k string) string {
	if t == nil {
		return ""
	}
	return t.m[k]
}

// Subset returns a new tags object with a subset of the keys.
func (t *Tags) Subset(keys []string) Tags {
	if len(keys) == 0 {
		return Tags{}
	}

	// If keys match existing keys, simply return this tagset.
	if keysMatch(t.m, keys) {
		return *t
	}

	// Otherwise create new tag set.
	m := make(map[string]string, len(keys))
	for _, k := range keys {
		m[k] = t.m[k]
	}
	return NewTags(m)
}

// Equals returns true if t equals other.
func (t *Tags) Equals(other *Tags) bool {
	if t == nil && other == nil {
		return true
	} else if t == nil || other == nil {
		return false
	}
	return t.id == other.id
}

// keysMatch returns true if m has exactly the same keys as listed in keys.
func keysMatch(m map[string]string, keys []string) bool {
	if len(keys) != len(m) {
		return false
	}

	for _, k := range keys {
		if _, ok := m[k]; !ok {
			return false
		}
	}

	return true
}

// encodeTags converts a map of strings to an identifier.
func encodeTags(m map[string]string) []byte {
	// Empty maps marshal to empty bytes.
	if len(m) == 0 {
		return nil
	}

	// Extract keys and determine final size.
	sz := (len(m) * 2) - 1 // separators
	keys := make([]string, 0, len(m))
	for k, v := range m {
		keys = append(keys, k)
		sz += len(k) + len(v)
	}
	sort.Strings(keys)

	// Generate marshaled bytes.
	b := make([]byte, sz)
	buf := b
	for _, k := range keys {
		copy(buf, k)
		buf[len(k)] = '\x00'
		buf = buf[len(k)+1:]
	}
	for i, k := range keys {
		v := m[k]
		copy(buf, v)
		if i < len(keys)-1 {
			buf[len(v)] = '\x00'
			buf = buf[len(v)+1:]
		}
	}
	return b
}

// decodeTags parses an identifier into a map of tags.
func decodeTags(id []byte) map[string]string {
	a := bytes.Split(id, []byte{'\x00'})

	// There must be an even number of segments.
	if len(a) > 0 && len(a)%2 == 1 {
		a = a[:len(a)-1]
	}

	// Return nil if there are no segments.
	if len(a) == 0 {
		return nil
	}
	mid := len(a) / 2

	// Decode key/value tags.
	m := make(map[string]string)
	for i := 0; i < mid; i++ {
		m[string(a[i])] = string(a[i+mid])
	}
	return m
}

func encodeAux(aux []interface{}) []*internal.Aux {
	pb := make([]*internal.Aux, len(aux))
	for i := range aux {
		switch v := aux[i].(type) {
		case float64:
			pb[i] = &internal.Aux{DataType: proto.Int32(Float), FloatValue: proto.Float64(v)}
		case *float64:
			pb[i] = &internal.Aux{DataType: proto.Int32(Float)}
		case int64:
			pb[i] = &internal.Aux{DataType: proto.Int32(Integer), IntegerValue: proto.Int64(v)}
		case *int64:
			pb[i] = &internal.Aux{DataType: proto.Int32(Integer)}
		case string:
			pb[i] = &internal.Aux{DataType: proto.Int32(String), StringValue: proto.String(v)}
		case *string:
			pb[i] = &internal.Aux{DataType: proto.Int32(String)}
		case bool:
			pb[i] = &internal.Aux{DataType: proto.Int32(Boolean), BooleanValue: proto.Bool(v)}
		case *bool:
			pb[i] = &internal.Aux{DataType: proto.Int32(Boolean)}
		default:
			pb[i] = &internal.Aux{DataType: proto.Int32(int32(Unknown))}
		}
	}
	return pb
}

func decodeAux(pb []*internal.Aux) []interface{} {
	if len(pb) == 0 {
		return nil
	}

	aux := make([]interface{}, len(pb))
	for i := range pb {
		switch pb[i].GetDataType() {
		case Float:
			if pb[i].FloatValue != nil {
				aux[i] = *pb[i].FloatValue
			} else {
				aux[i] = (*float64)(nil)
			}
		case Integer:
			if pb[i].IntegerValue != nil {
				aux[i] = *pb[i].IntegerValue
			} else {
				aux[i] = (*int64)(nil)
			}
		case String:
			if pb[i].StringValue != nil {
				aux[i] = *pb[i].StringValue
			} else {
				aux[i] = (*string)(nil)
			}
		case Boolean:
			if pb[i].BooleanValue != nil {
				aux[i] = *pb[i].BooleanValue
			} else {
				aux[i] = (*bool)(nil)
			}
		default:
			aux[i] = nil
		}
	}
	return aux
}

// NewPointDecoder decodes generic points from a reader.
type PointDecoder struct {
	r     io.Reader
	stats IteratorStats
}

// NewPointDecoder returns a new instance of PointDecoder that reads from r.
func NewPointDecoder(r io.Reader) *PointDecoder {
	return &PointDecoder{r: r}
}

// Stats returns iterator stats embedded within the stream.
func (dec *PointDecoder) Stats() IteratorStats { return dec.stats }

// DecodePoint reads from the underlying reader and unmarshals into p.
func (dec *PointDecoder) DecodePoint(p *Point) error {
	for {
		// Read length.
		var sz uint32
		if err := binary.Read(dec.r, binary.BigEndian, &sz); err != nil {
			return err
		}

		// Read point data.
		buf := make([]byte, sz)
		if _, err := io.ReadFull(dec.r, buf); err != nil {
			return err
		}

		// Unmarshal into point.
		var pb internal.Point
		if err := proto.Unmarshal(buf, &pb); err != nil {
			return err
		}

		// If the point contains stats then read stats and retry.
		if pb.Stats != nil {
			dec.stats = decodeIteratorStats(pb.Stats)
			continue
		}

		if pb.IntegerValue != nil {
			*p = decodeIntegerPoint(&pb)
		} else if pb.StringValue != nil {
			*p = decodeStringPoint(&pb)
		} else if pb.BooleanValue != nil {
			*p = decodeBooleanPoint(&pb)
		} else {
			*p = decodeFloatPoint(&pb)
		}

		return nil
	}
}
<|endoftext|>


<commit_before>//
// Copyright (c) 2016 Nicolas Martyanoff <khaelin@gmail.com>
//
// Permission to use, copy, modify, and distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
// ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
// ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
// OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

package imapc

import (
	"bytes"
	"fmt"
	"strconv"
)

type Response interface {
	fmt.GoStringer
	Read(*Stream) error
}

func ReadResponse(s *Stream) (Response, error) {
	prefix, err := s.Peek(1)
	if err != nil {
		return nil, err
	}

	var readFn func(*Stream) (Response, error)

	if prefix[0] == '+' {
		// Command continuation
		readFn = ReadResponseContinuation
	} else if prefix[0] == '*' {
		// Server data or untagged status response
		// The only untagged status response is BYE.
		const byePrefix = "* BYE "
		prefix, err := s.Peek(len(byePrefix))
		if err != nil {
			return nil, err
		}

		if bytes.Equal(prefix, []byte(byePrefix)) {
			// Untagged status response
			readFn = ReadResponseStatus
		} else {
			// Server data
			readFn = ReadResponseData
		}
	} else {
		// Tagged status response
		readFn = ReadResponseStatus
	}

	r, err := readFn(s)
	if err != nil {
		return nil, err
	}

	if err := r.Read(s); err != nil {
		return nil, err
	}

	return r, nil
}

// ---------------------------------------------------------------------------
//  Status responses
// ---------------------------------------------------------------------------
func ReadResponseStatus(s *Stream) (Response, error) {
	// Read the tag
	tag, err := s.ReadUntilByteAndSkip(' ')
	if err != nil {
		return nil, err
	}

	// Read the response name
	name, err := s.ReadUntilByteAndSkip(' ')
	if err != nil {
		return nil, err
	}

	r := &ResponseStatus{
		Tag:          string(tag),
		ResponseName: string(name),
	}

	switch string(name) {
	case "OK":
		r.Response = &ResponseOk{}
	case "NO":
		r.Response = &ResponseNo{}
	case "BAD":
		r.Response = &ResponseBad{}
	case "BYE":
		r.Response = &ResponseBye{}
	default:
		return nil, fmt.Errorf("unknown response %q", name)
	}

	return r, nil
}

type ResponseStatus struct {
	Tag          string
	ResponseName string
	Response     Response
}

func (r *ResponseStatus) GoString() string {
	return fmt.Sprintf("#<response-status %#v>", r.Response)
}

func (r *ResponseStatus) IsOk() bool {
	_, ok := r.Response.(*ResponseOk)
	return ok
}

func (r *ResponseStatus) Read(s *Stream) error {
	return r.Response.Read(s)
}

// OK
type ResponseOk struct {
	Text *ResponseText
}

func (r *ResponseOk) GoString() string {
	return fmt.Sprintf("#<response-ok %#v>", r.Text)
}

func (r *ResponseOk) Read(s *Stream) error {
	r.Text = &ResponseText{}
	return r.Text.Read(s)
}

// NO
type ResponseNo struct {
	Text *ResponseText
}

func (r *ResponseNo) GoString() string {
	return fmt.Sprintf("#<response-no %#v>", r.Text)
}

func (r *ResponseNo) Read(s *Stream) error {
	r.Text = &ResponseText{}
	return r.Text.Read(s)
}

// BAD
type ResponseBad struct {
	Text *ResponseText
}

func (r *ResponseBad) GoString() string {
	return fmt.Sprintf("#<response-bad %#v>", r.Text)
}

func (r *ResponseBad) Read(s *Stream) error {
	r.Text = &ResponseText{}
	return r.Text.Read(s)
}

// ---------------------------------------------------------------------------
//  Server data responses
// ---------------------------------------------------------------------------
func ReadResponseData(s *Stream) (Response, error) {
	// Skip "* "
	if err := s.Skip(2); err != nil {
		return nil, err
	}

	// Read either the response tag or a number
	data, err := s.ReadUntilByteAndSkip(' ')
	if err != nil {
		return nil, err
	}

	var r Response

	if ByteStringAll(data, IsDigitChar) {
		n, err := strconv.ParseUint(string(data), 10, 32)
		if err != nil {
			return nil, err
		}
		count := uint32(n)

		rest, err := s.ReadUntil([]byte("\r\n"))
		if err != nil {
			return nil, err
		}
		tag := string(rest)

		switch tag {
		case "EXISTS":
			r = &ResponseExists{Count: count}
		case "RECENT":
			r = &ResponseRecent{Count: count}
		default:
			return nil, fmt.Errorf("unknown response %q", tag)
		}
	} else {
		tag := string(data)

		switch tag {
		case "OK":
			r = &ResponseOk{}
		case "PREAUTH":
			r = &ResponsePreAuth{}
		case "BYE":
			r = &ResponseBye{}
		case "CAPABILITY":
			r = &ResponseCapability{}
		case "LIST":
			r = &ResponseList{}
		case "LSUB":
			r = &ResponseLSub{}
		case "FLAGS":
			r = &ResponseFlags{}
		default:
			return nil, fmt.Errorf("unknown response %q", tag)
		}
	}

	return r, nil
}

// PREAUTH
type ResponsePreAuth struct {
	Text *ResponseText
}

func (r *ResponsePreAuth) GoString() string {
	return fmt.Sprintf("#<response-pre-auth %#v>", r.Text)
}

func (r *ResponsePreAuth) Read(s *Stream) error {
	r.Text = &ResponseText{}
	return r.Text.Read(s)
}

// BYE
type ResponseBye struct {
	Text *ResponseText
}

func (r *ResponseBye) GoString() string {
	return fmt.Sprintf("#<response-bye %#v>", r.Text)
}

func (r *ResponseBye) Read(s *Stream) error {
	r.Text = &ResponseText{}
	return r.Text.Read(s)
}

// CAPABILITY
type ResponseCapability struct {
	Caps []string
}

func (r *ResponseCapability) GoString() string {
	return fmt.Sprintf("#<response-capability %v>", r.Caps)
}

func (r *ResponseCapability) Read(s *Stream) error {
	data, err := s.ReadUntilAndSkip([]byte("\r\n"))
	if err != nil {
		return err
	}

	parts := bytes.Split(data, []byte{' '})

	r.Caps = make([]string, len(parts))
	for i, cap := range parts {
		r.Caps[i] = string(cap)
	}

	return nil
}

// LIST
type ResponseList MailboxList

func (r *ResponseList) GoString() string {
	return fmt.Sprintf("#<response-list %q %v>", r.Name, r.Flags)
}

func (r *ResponseList) Read(s *Stream) error {
	mbox, err := s.ReadIMAPMailboxList()
	if err != nil {
		return err
	}
	*r = ResponseList(*mbox)

	if ok, err := s.SkipBytes([]byte("\r\n")); err != nil {
		return err
	} else if !ok {
		return fmt.Errorf("invalid character after mailbox list")
	}

	return nil
}

// LSUB
type ResponseLSub MailboxList

func (r *ResponseLSub) GoString() string {
	return fmt.Sprintf("#<response-lsub %q %v>", r.Name, r.Flags)
}

func (r *ResponseLSub) Read(s *Stream) error {
	mbox, err := s.ReadIMAPMailboxList()
	if err != nil {
		return err
	}
	*r = ResponseLSub(*mbox)

	if ok, err := s.SkipBytes([]byte("\r\n")); err != nil {
		return err
	} else if !ok {
		return fmt.Errorf("invalid character after mailbox list")
	}

	return nil
}

// FLAGS
type ResponseFlags struct {
	Flags []string
}

func (r *ResponseFlags) GoString() string {
	return fmt.Sprintf("#<response-flags %v>", r.Flags)
}

func (r *ResponseFlags) Read(s *Stream) error {
	// Flags
	flags, err := s.ReadIMAPFlagList()
	if err != nil {
		return err
	}

	r.Flags = flags

	// End
	if ok, err := s.SkipBytes([]byte("\r\n")); err != nil {
		return err
	} else if !ok {
		return fmt.Errorf("invalid character after flag list")
	}

	return nil
}

// EXISTS
type ResponseExists struct {
	Count uint32
}

func (r *ResponseExists) GoString() string {
	return fmt.Sprintf("#<response-exists %d>", r.Count)
}

func (r *ResponseExists) Read(s *Stream) error {
	// End
	if ok, err := s.SkipBytes([]byte("\r\n")); err != nil {
		return err
	} else if !ok {
		return fmt.Errorf("invalid character after EXISTS")
	}

	return nil
}

// RECENT
type ResponseRecent struct {
	Count uint32
}

func (r *ResponseRecent) GoString() string {
	return fmt.Sprintf("#<response-recent %d>", r.Count)
}

func (r *ResponseRecent) Read(s *Stream) error {
	// End
	if ok, err := s.SkipBytes([]byte("\r\n")); err != nil {
		return err
	} else if !ok {
		return fmt.Errorf("invalid character after RECENT")
	}

	return nil
}

// ---------------------------------------------------------------------------
//  Command continuation responses
// ---------------------------------------------------------------------------
func ReadResponseContinuation(s *Stream) (Response, error) {
	// Skip "+ "
	if err := s.Skip(2); err != nil {
		return nil, err
	}

	return &ResponseContinuation{}, nil
}

type ResponseContinuation struct {
	Text string
}

func (r *ResponseContinuation) GoString() string {
	return fmt.Sprintf("#<response-continuation %q>", r.Text)
}

func (r *ResponseContinuation) Read(s *Stream) error {
	text, err := s.ReadUntilAndSkip([]byte("\r\n"))
	if err != nil {
		return err
	}
	r.Text = string(text)

	return nil
}

// ---------------------------------------------------------------------------
//  Response text
// ---------------------------------------------------------------------------
type ResponseText struct {
	Text string

	// Optional
	Code       string
	CodeString string
	CodeData   interface{}
}

func (r *ResponseText) Read(s *Stream) error {
	found, err := s.SkipByte('[')
	if err != nil {
		return err
	} else if found {
		// Response text code
		codeBytes, err := s.ReadUntilAndSkip([]byte("] "))
		if err != nil {
			return err
		}
		r.CodeString = string(codeBytes)

		var codeData []byte

		idx := bytes.IndexByte(codeBytes, ' ')
		if idx == -1 {
			r.Code = string(codeBytes)
			codeData = []byte{}
		} else {
			r.Code = string(codeBytes[0:idx])
			codeData = codeBytes[idx+1:]
		}

		if err := r.ReadCodeData(codeData); err != nil {
			return fmt.Errorf("invalid %s data: %v", r.Code, err)
		}
	}

	// Response text
	text, err := s.ReadUntilAndSkip([]byte("\r\n"))
	if err != nil {
		return err
	}

	r.Text = string(text)
	return nil
}

func (r *ResponseText) ReadCodeData(data []byte) error {
	s := NewStream(bytes.NewReader(data))

	switch r.Code {
	case "CAPABILITY":
		capsData, err := s.ReadAll()
		if err != nil {
			return err
		}

		parts := bytes.Split(capsData, []byte{' '})
		caps := make([]string, len(parts))
		for i, cap := range parts {
			caps[i] = string(cap)
		}

		r.CodeData = caps

	case "HIGHESTMODSEQ":
		fallthrough
	case "UIDNEXT":
		fallthrough
	case "UIDVALIDITY":
		fallthrough
	case "UNSEEN":
		n, err := s.ReadIMAPNumber()
		if err != nil {
			return err
		} else if n == 0 {
			return fmt.Errorf("invalid zero value")
		}

		r.CodeData = n

	case "PERMANENTFLAGS":
		flags, err := s.ReadIMAPFlagList()
		if err != nil {
			return err
		}

		r.CodeData = flags

	default:
		codeData, err := s.ReadAll()
		if err != nil {
			return err
		}

		r.CodeData = codeData
	}

	empty, err := s.IsEmpty()
	if err != nil {
		return err
	} else if !empty {
		return fmt.Errorf("invalid trailing data")
	}

	return nil
}

func (r *ResponseText) GoString() string {
	if r.Code == "" {
		return fmt.Sprintf("#<response-text %q>", r.Text)
	} else {
		return fmt.Sprintf("#<response-text %s %v %q>",
			r.Code, r.CodeData, r.Text)
	}
}
<commit_msg>fix data response parsing<commit_after>//
// Copyright (c) 2016 Nicolas Martyanoff <khaelin@gmail.com>
//
// Permission to use, copy, modify, and distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
// ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
// ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
// OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

package imapc

import (
	"bytes"
	"fmt"
	"strconv"
)

type Response interface {
	fmt.GoStringer
	Read(*Stream) error
}

func ReadResponse(s *Stream) (Response, error) {
	prefix, err := s.Peek(1)
	if err != nil {
		return nil, err
	}

	var readFn func(*Stream) (Response, error)

	if prefix[0] == '+' {
		// Command continuation
		readFn = ReadResponseContinuation
	} else if prefix[0] == '*' {
		// Server data or untagged status response
		// The only untagged status response is BYE.
		const byePrefix = "* BYE "
		prefix, err := s.Peek(len(byePrefix))
		if err != nil {
			return nil, err
		}

		if bytes.Equal(prefix, []byte(byePrefix)) {
			// Untagged status response
			readFn = ReadResponseStatus
		} else {
			// Server data
			readFn = ReadResponseData
		}
	} else {
		// Tagged status response
		readFn = ReadResponseStatus
	}

	r, err := readFn(s)
	if err != nil {
		return nil, err
	}

	if err := r.Read(s); err != nil {
		return nil, err
	}

	return r, nil
}

// ---------------------------------------------------------------------------
//  Status responses
// ---------------------------------------------------------------------------
func ReadResponseStatus(s *Stream) (Response, error) {
	// Read the tag
	tag, err := s.ReadUntilByteAndSkip(' ')
	if err != nil {
		return nil, err
	}

	// Read the response name
	name, err := s.ReadUntilByteAndSkip(' ')
	if err != nil {
		return nil, err
	}

	r := &ResponseStatus{
		Tag:          string(tag),
		ResponseName: string(name),
	}

	switch string(name) {
	case "OK":
		r.Response = &ResponseOk{}
	case "NO":
		r.Response = &ResponseNo{}
	case "BAD":
		r.Response = &ResponseBad{}
	case "BYE":
		r.Response = &ResponseBye{}
	default:
		return nil, fmt.Errorf("unknown response %q", name)
	}

	return r, nil
}

type ResponseStatus struct {
	Tag          string
	ResponseName string
	Response     Response
}

func (r *ResponseStatus) GoString() string {
	return fmt.Sprintf("#<response-status %#v>", r.Response)
}

func (r *ResponseStatus) IsOk() bool {
	_, ok := r.Response.(*ResponseOk)
	return ok
}

func (r *ResponseStatus) Read(s *Stream) error {
	return r.Response.Read(s)
}

// OK
type ResponseOk struct {
	Text *ResponseText
}

func (r *ResponseOk) GoString() string {
	return fmt.Sprintf("#<response-ok %#v>", r.Text)
}

func (r *ResponseOk) Read(s *Stream) error {
	r.Text = &ResponseText{}
	return r.Text.Read(s)
}

// NO
type ResponseNo struct {
	Text *ResponseText
}

func (r *ResponseNo) GoString() string {
	return fmt.Sprintf("#<response-no %#v>", r.Text)
}

func (r *ResponseNo) Read(s *Stream) error {
	r.Text = &ResponseText{}
	return r.Text.Read(s)
}

// BAD
type ResponseBad struct {
	Text *ResponseText
}

func (r *ResponseBad) GoString() string {
	return fmt.Sprintf("#<response-bad %#v>", r.Text)
}

func (r *ResponseBad) Read(s *Stream) error {
	r.Text = &ResponseText{}
	return r.Text.Read(s)
}

// ---------------------------------------------------------------------------
//  Server data responses
// ---------------------------------------------------------------------------
func ReadResponseData(s *Stream) (Response, error) {
	// Skip "* "
	if err := s.Skip(2); err != nil {
		return nil, err
	}

	// Read either the response tag or a number
	data, err := s.ReadWhile(func(b byte) bool {
		return (b >= 'A' && b <= 'Z') || IsDigitChar(b)
	})
	if err != nil {
		return nil, err
	}

	if _, err := s.SkipByte(' '); err != nil {
		return nil, err
	}

	var r Response

	if ByteStringAll(data, IsDigitChar) {
		n, err := strconv.ParseUint(string(data), 10, 32)
		if err != nil {
			return nil, err
		}
		count := uint32(n)

		rest, err := s.ReadUntil([]byte("\r\n"))
		if err != nil {
			return nil, err
		}
		tag := string(rest)

		switch tag {
		case "EXISTS":
			r = &ResponseExists{Count: count}
		case "RECENT":
			r = &ResponseRecent{Count: count}
		default:
			return nil, fmt.Errorf("unknown response %q", tag)
		}
	} else {
		tag := string(data)

		switch tag {
		case "OK":
			r = &ResponseOk{}
		case "PREAUTH":
			r = &ResponsePreAuth{}
		case "BYE":
			r = &ResponseBye{}
		case "CAPABILITY":
			r = &ResponseCapability{}
		case "LIST":
			r = &ResponseList{}
		case "LSUB":
			r = &ResponseLSub{}
		case "FLAGS":
			r = &ResponseFlags{}
		default:
			return nil, fmt.Errorf("unknown response %q", tag)
		}
	}

	return r, nil
}

// PREAUTH
type ResponsePreAuth struct {
	Text *ResponseText
}

func (r *ResponsePreAuth) GoString() string {
	return fmt.Sprintf("#<response-pre-auth %#v>", r.Text)
}

func (r *ResponsePreAuth) Read(s *Stream) error {
	r.Text = &ResponseText{}
	return r.Text.Read(s)
}

// BYE
type ResponseBye struct {
	Text *ResponseText
}

func (r *ResponseBye) GoString() string {
	return fmt.Sprintf("#<response-bye %#v>", r.Text)
}

func (r *ResponseBye) Read(s *Stream) error {
	r.Text = &ResponseText{}
	return r.Text.Read(s)
}

// CAPABILITY
type ResponseCapability struct {
	Caps []string
}

func (r *ResponseCapability) GoString() string {
	return fmt.Sprintf("#<response-capability %v>", r.Caps)
}

func (r *ResponseCapability) Read(s *Stream) error {
	data, err := s.ReadUntilAndSkip([]byte("\r\n"))
	if err != nil {
		return err
	}

	parts := bytes.Split(data, []byte{' '})

	r.Caps = make([]string, len(parts))
	for i, cap := range parts {
		r.Caps[i] = string(cap)
	}

	return nil
}

// LIST
type ResponseList MailboxList

func (r *ResponseList) GoString() string {
	return fmt.Sprintf("#<response-list %q %v>", r.Name, r.Flags)
}

func (r *ResponseList) Read(s *Stream) error {
	mbox, err := s.ReadIMAPMailboxList()
	if err != nil {
		return err
	}
	*r = ResponseList(*mbox)

	if ok, err := s.SkipBytes([]byte("\r\n")); err != nil {
		return err
	} else if !ok {
		return fmt.Errorf("invalid character after mailbox list")
	}

	return nil
}

// LSUB
type ResponseLSub MailboxList

func (r *ResponseLSub) GoString() string {
	return fmt.Sprintf("#<response-lsub %q %v>", r.Name, r.Flags)
}

func (r *ResponseLSub) Read(s *Stream) error {
	mbox, err := s.ReadIMAPMailboxList()
	if err != nil {
		return err
	}
	*r = ResponseLSub(*mbox)

	if ok, err := s.SkipBytes([]byte("\r\n")); err != nil {
		return err
	} else if !ok {
		return fmt.Errorf("invalid character after mailbox list")
	}

	return nil
}

// FLAGS
type ResponseFlags struct {
	Flags []string
}

func (r *ResponseFlags) GoString() string {
	return fmt.Sprintf("#<response-flags %v>", r.Flags)
}

func (r *ResponseFlags) Read(s *Stream) error {
	// Flags
	flags, err := s.ReadIMAPFlagList()
	if err != nil {
		return err
	}

	r.Flags = flags

	// End
	if ok, err := s.SkipBytes([]byte("\r\n")); err != nil {
		return err
	} else if !ok {
		return fmt.Errorf("invalid character after flag list")
	}

	return nil
}

// EXISTS
type ResponseExists struct {
	Count uint32
}

func (r *ResponseExists) GoString() string {
	return fmt.Sprintf("#<response-exists %d>", r.Count)
}

func (r *ResponseExists) Read(s *Stream) error {
	// End
	if ok, err := s.SkipBytes([]byte("\r\n")); err != nil {
		return err
	} else if !ok {
		return fmt.Errorf("invalid character after EXISTS")
	}

	return nil
}

// RECENT
type ResponseRecent struct {
	Count uint32
}

func (r *ResponseRecent) GoString() string {
	return fmt.Sprintf("#<response-recent %d>", r.Count)
}

func (r *ResponseRecent) Read(s *Stream) error {
	// End
	if ok, err := s.SkipBytes([]byte("\r\n")); err != nil {
		return err
	} else if !ok {
		return fmt.Errorf("invalid character after RECENT")
	}

	return nil
}

// ---------------------------------------------------------------------------
//  Command continuation responses
// ---------------------------------------------------------------------------
func ReadResponseContinuation(s *Stream) (Response, error) {
	// Skip "+ "
	if err := s.Skip(2); err != nil {
		return nil, err
	}

	return &ResponseContinuation{}, nil
}

type ResponseContinuation struct {
	Text string
}

func (r *ResponseContinuation) GoString() string {
	return fmt.Sprintf("#<response-continuation %q>", r.Text)
}

func (r *ResponseContinuation) Read(s *Stream) error {
	text, err := s.ReadUntilAndSkip([]byte("\r\n"))
	if err != nil {
		return err
	}
	r.Text = string(text)

	return nil
}

// ---------------------------------------------------------------------------
//  Response text
// ---------------------------------------------------------------------------
type ResponseText struct {
	Text string

	// Optional
	Code       string
	CodeString string
	CodeData   interface{}
}

func (r *ResponseText) Read(s *Stream) error {
	found, err := s.SkipByte('[')
	if err != nil {
		return err
	} else if found {
		// Response text code
		codeBytes, err := s.ReadUntilAndSkip([]byte("] "))
		if err != nil {
			return err
		}
		r.CodeString = string(codeBytes)

		var codeData []byte

		idx := bytes.IndexByte(codeBytes, ' ')
		if idx == -1 {
			r.Code = string(codeBytes)
			codeData = []byte{}
		} else {
			r.Code = string(codeBytes[0:idx])
			codeData = codeBytes[idx+1:]
		}

		if err := r.ReadCodeData(codeData); err != nil {
			return fmt.Errorf("invalid %s data: %v", r.Code, err)
		}
	}

	// Response text
	text, err := s.ReadUntilAndSkip([]byte("\r\n"))
	if err != nil {
		return err
	}

	r.Text = string(text)
	return nil
}

func (r *ResponseText) ReadCodeData(data []byte) error {
	s := NewStream(bytes.NewReader(data))

	switch r.Code {
	case "CAPABILITY":
		capsData, err := s.ReadAll()
		if err != nil {
			return err
		}

		parts := bytes.Split(capsData, []byte{' '})
		caps := make([]string, len(parts))
		for i, cap := range parts {
			caps[i] = string(cap)
		}

		r.CodeData = caps

	case "HIGHESTMODSEQ":
		fallthrough
	case "UIDNEXT":
		fallthrough
	case "UIDVALIDITY":
		fallthrough
	case "UNSEEN":
		n, err := s.ReadIMAPNumber()
		if err != nil {
			return err
		} else if n == 0 {
			return fmt.Errorf("invalid zero value")
		}

		r.CodeData = n

	case "PERMANENTFLAGS":
		flags, err := s.ReadIMAPFlagList()
		if err != nil {
			return err
		}

		r.CodeData = flags

	default:
		codeData, err := s.ReadAll()
		if err != nil {
			return err
		}

		r.CodeData = codeData
	}

	empty, err := s.IsEmpty()
	if err != nil {
		return err
	} else if !empty {
		return fmt.Errorf("invalid trailing data")
	}

	return nil
}

func (r *ResponseText) GoString() string {
	if r.Code == "" {
		return fmt.Sprintf("#<response-text %q>", r.Text)
	} else {
		return fmt.Sprintf("#<response-text %s %v %q>",
			r.Code, r.CodeData, r.Text)
	}
}
<|endoftext|>