<commit_before>//
// message.go
//
package srnd

import (
  "bufio"
  "encoding/base64"
  "fmt"
  "io"
  "log"
  "mime"
  "mime/multipart"
  "net/textproto"
  "strings"
  "time"
)

type ArticleHeaders map[string][]string

func (self ArticleHeaders) Has(key string) bool {
  _, ok := self[key]
  return ok
}

func (self ArticleHeaders) Set(key, val string) {
  self[key] = []string{val}
}

func (self ArticleHeaders) Get(key, fallback string) string {
  val , ok := self[key]
  if ok {
    return val[0]
  } else {
    return fallback
  }
}

type NNTPMessage interface {
  // this message's messsge id
  MessageID() string
  // the parent message's messageid if it's specified
  Reference() string
  // the newsgroup this post is in
  Newsgroup() string
  // the name of the poster
  Name() string
  // any email address associated with the post
  Email() string
  // the subject of the post
  Subject() string
  // when this was posted
  Posted() int64
  // the path header
  Path() string
  // append something to path
  // return message with new path
  AppendPath(part string) NNTPMessage
  // the type of this message usually a mimetype
  ContentType() string
  // was this post a sage?
  Sage() bool
  // was this post a root post?
  OP() bool
  // all attachments
  Attachments() []NNTPAttachment
  // all headers
  Headers() ArticleHeaders
  // write out body
  WriteBody(wr io.Writer) error
  // attach a file
  Attach(att NNTPAttachment) NNTPMessage
  // get the plaintext message if it exists
  Message() string
  // pack the whole message and prepare for write
  Pack()
  // get the inner nntp article that is signed and valid, returns nil if not present or invalid
  Signed() NNTPMessage
}

type MessageReader interface {
  // read a message from a reader
  ReadMessage(r io.Reader) (NNTPMessage, error)
}

type MessageWriter interface {
  // write a message to a writer
  WriteMessage(nntp NNTPMessage, wr io.Writer) error
}

type nntpArticle struct {
  headers ArticleHeaders
  boundary string
  message nntpAttachment
  attachments []NNTPAttachment
  signedPart nntpAttachment
}

// create a simple plaintext nntp message
func newPlaintextArticle(message, email, subject, name, instance, newsgroup string) NNTPMessage {
  nntp := nntpArticle{
    headers: make(ArticleHeaders),
  }
  nntp.headers.Set("From", fmt.Sprintf("%s <%s>", name, email))
  nntp.headers.Set("Subject", subject)
  if isSage(subject) {
    nntp.headers.Set("X-Sage", "1")
  }
  nntp.headers.Set("Path", instance)
  nntp.headers.Set("Message-ID", genMessageID(instance))
  // posted now
  nntp.headers.Set("Date", timeNowStr())
  nntp.headers.Set("Newsgroups", newsgroup)
  nntp.message = createPlaintextAttachment(message)
  nntp.Pack()
  return nntp
}

func (self nntpArticle) Signed() NNTPMessage {
  if self.signedPart.body.Len() > 0 {
    msg, err := read_message(&self.signedPart.body)
    if err == nil {
      return msg
    }
    log.Println("failed to load signed message", err)
  }
  return nil
}

func (self nntpArticle) MessageID() string {
  return self.headers.Get("Message-ID", self.headers.Get("Messageid", self.headers.Get("MessageID", self.headers.Get("Message-Id", ""))))
}

func (self nntpArticle) Pack() {
  if len(self.attachments) > 0 {
    if len(self.boundary) == 0 {
      // we have no boundry, set it
      self.boundary = randStr(24)
      // set headers
      self.headers.Set("Mime-Version", "1.0")
      self.headers.Set("Content-Type", fmt.Sprintf("multipart/mixed; boundary=%s", self.boundary))
    }
  } else {
    self.headers.Set("Content-Type", "text/plain; charset=utf-8")
  }

}

func (self nntpArticle) Reference() string {
  return self.headers.Get("Reference", self.headers.Get("References",""))
}

func (self nntpArticle) Newsgroup() string {
  return self.headers.Get("Newsgroups", "")
}

func (self nntpArticle) Name() string {
  from := self.headers.Get("From", "anonymous <a@no.n>")
  idx := strings.Index(from, " ")
  if idx > 1 {
    return from[:idx]
  }
  return "[Invalid From header]"
}

func (self nntpArticle) Email() string {
  from := self.headers.Get("From", "anonymous <a@no.n>")
  idx := strings.Index(from, " ")
  if idx > 1 {
    idx_1 := strings.Index(from[:idx], "<")
    idx_2 := strings.Index(from[:idx], ">")
    if idx_2 > 0 && idx_1 > 0 && idx_2 > idx_1 {
      return from[1+idx+idx_1:idx+idx_2]
    }
  }
  return "[Invalid From header]"

}

func (self nntpArticle) Subject() string {
  return self.headers.Get("Subject", "")
}

func (self nntpArticle) Posted() int64 {
  posted := self.headers.Get("Date", "")
  t, err := time.Parse(time.RFC1123Z, posted)
  if err == nil {
    return t.Unix()
  }
  return 0
}

func (self nntpArticle) Message() string {
  return self.message.body.String()
}

func (self nntpArticle) Path() string {
  return self.headers.Get("Path", "unspecified")
}

func (self nntpArticle) Headers() ArticleHeaders {
  return self.headers
}

func (self nntpArticle) AppendPath(part string) NNTPMessage {
  if self.headers.Has("Path") {
    self.headers.Set("Path", part + "!" + self.Path())
  } else {
    self.headers.Set("Path", part)
  }
  return self
}
func (self nntpArticle) ContentType() string {
  // assumes text/plain if unspecified
  return self.headers.Get("Content-Type", "text/plain")
}

func (self nntpArticle) Sage() bool {
  return self.headers.Has("X-Sage")
}

func (self nntpArticle) OP() bool {
  return self.headers.Get("Reference", self.headers.Get("References", "")) == ""
}

func (self nntpArticle) Attachments() []NNTPAttachment {
  return self.attachments
}

func (self nntpArticle) Attach(att NNTPAttachment) NNTPMessage {
  log.Println("attaching file", att.Filename())
  self.attachments = append(self.attachments, att)
  return self
}

func (self nntpArticle) WriteBody(wr io.Writer) (err error) {
  // this is a signed message, don't treat it special
  if self.signedPart.body.Len() > 0 {
    r := bufio.NewReader(&self.signedPart.body)
    w := bufio.NewWriter(wr)
    var line []byte
    for {
      // convert line endings :\
      line, err = r.ReadBytes('\n')
      if err == nil {
        w.Write(line[:len(line)-2])
        w.WriteByte(10)
      } else {
        return
      }
    }
  }
  if len(self.attachments) == 0 {
    // write plaintext and be done
    _, err = self.message.body.WriteTo(wr)
    return
  }
  content_type := self.ContentType()
  _, params, err := mime.ParseMediaType(content_type)
  if err != nil {
    log.Println("failed to parse media type", err)
    return err
  }

  boundary, ok := params["boundary"]
  if ok {
    w := multipart.NewWriter(wr)

    err = w.SetBoundary(boundary)
    if err == nil {
      attachments := []NNTPAttachment{self.message}
      attachments = append(attachments, self.attachments...)
      log.Println("writing nntp body")
      for _ , att := range(attachments) {
        hdr := att.Header()
        if hdr == nil {
          hdr = make(textproto.MIMEHeader)
        }
        hdr.Set("Content-Transfer-Encoding", "base64")
        part, err := w.CreatePart(hdr)
        enc := base64.NewEncoder(base64.StdEncoding, part)
        err = att.WriteTo(enc)
        enc.Close()

        if err != nil {
          break
        }
      }
    }
    if err != nil {
      log.Println(err)
    }
    err = w.Close()
  } else {
    _, err = self.message.body.WriteTo(wr)
  }
  return err
}
<commit_msg>change how article writing works<commit_after>//
// message.go
//
package srnd

import (
  "bufio"
  "encoding/base64"
  "fmt"
  "io"
  "log"
  "mime"
  "mime/multipart"
  "net/textproto"
  "strings"
  "time"
)

type ArticleHeaders map[string][]string

func (self ArticleHeaders) Has(key string) bool {
  _, ok := self[key]
  return ok
}

func (self ArticleHeaders) Set(key, val string) {
  self[key] = []string{val}
}

func (self ArticleHeaders) Get(key, fallback string) string {
  val , ok := self[key]
  if ok {
    return val[0]
  } else {
    return fallback
  }
}

type NNTPMessage interface {
  // this message's messsge id
  MessageID() string
  // the parent message's messageid if it's specified
  Reference() string
  // the newsgroup this post is in
  Newsgroup() string
  // the name of the poster
  Name() string
  // any email address associated with the post
  Email() string
  // the subject of the post
  Subject() string
  // when this was posted
  Posted() int64
  // the path header
  Path() string
  // append something to path
  // return message with new path
  AppendPath(part string) NNTPMessage
  // the type of this message usually a mimetype
  ContentType() string
  // was this post a sage?
  Sage() bool
  // was this post a root post?
  OP() bool
  // all attachments
  Attachments() []NNTPAttachment
  // all headers
  Headers() ArticleHeaders
  // write out body
  WriteBody(wr io.Writer) error
  // attach a file
  Attach(att NNTPAttachment) NNTPMessage
  // get the plaintext message if it exists
  Message() string
  // pack the whole message and prepare for write
  Pack()
  // get the inner nntp article that is signed and valid, returns nil if not present or invalid
  Signed() NNTPMessage
}

type MessageReader interface {
  // read a message from a reader
  ReadMessage(r io.Reader) (NNTPMessage, error)
}

type MessageWriter interface {
  // write a message to a writer
  WriteMessage(nntp NNTPMessage, wr io.Writer) error
}

type nntpArticle struct {
  headers ArticleHeaders
  boundary string
  message nntpAttachment
  attachments []NNTPAttachment
  signedPart nntpAttachment
}

// create a simple plaintext nntp message
func newPlaintextArticle(message, email, subject, name, instance, newsgroup string) NNTPMessage {
  nntp := nntpArticle{
    headers: make(ArticleHeaders),
  }
  nntp.headers.Set("From", fmt.Sprintf("%s <%s>", name, email))
  nntp.headers.Set("Subject", subject)
  if isSage(subject) {
    nntp.headers.Set("X-Sage", "1")
  }
  nntp.headers.Set("Path", instance)
  nntp.headers.Set("Message-ID", genMessageID(instance))
  // posted now
  nntp.headers.Set("Date", timeNowStr())
  nntp.headers.Set("Newsgroups", newsgroup)
  nntp.message = createPlaintextAttachment(message)
  nntp.Pack()
  return nntp
}

func (self nntpArticle) Signed() NNTPMessage {
  if self.signedPart.body.Len() > 0 {
    msg, err := read_message(&self.signedPart.body)
    if err == nil {
      return msg
    }
    log.Println("failed to load signed message", err)
  }
  return nil
}

func (self nntpArticle) MessageID() string {
  return self.headers.Get("Message-ID", self.headers.Get("Messageid", self.headers.Get("MessageID", self.headers.Get("Message-Id", ""))))
}

func (self nntpArticle) Pack() {
  if len(self.attachments) > 0 {
    if len(self.boundary) == 0 {
      // we have no boundry, set it
      self.boundary = randStr(24)
      // set headers
      self.headers.Set("Mime-Version", "1.0")
      self.headers.Set("Content-Type", fmt.Sprintf("multipart/mixed; boundary=%s", self.boundary))
    }
  } else if self.signedPart.body.Len() == 0 {
    self.headers.Set("Content-Type", "text/plain; charset=utf-8")
  }

}

func (self nntpArticle) Reference() string {
  return self.headers.Get("Reference", self.headers.Get("References",""))
}

func (self nntpArticle) Newsgroup() string {
  return self.headers.Get("Newsgroups", "")
}

func (self nntpArticle) Name() string {
  from := self.headers.Get("From", "anonymous <a@no.n>")
  idx := strings.Index(from, " ")
  if idx > 1 {
    return from[:idx]
  }
  return "[Invalid From header]"
}

func (self nntpArticle) Email() string {
  from := self.headers.Get("From", "anonymous <a@no.n>")
  idx := strings.Index(from, " ")
  if idx > 1 {
    idx_1 := strings.Index(from[:idx], "<")
    idx_2 := strings.Index(from[:idx], ">")
    if idx_2 > 0 && idx_1 > 0 && idx_2 > idx_1 {
      return from[1+idx+idx_1:idx+idx_2]
    }
  }
  return "[Invalid From header]"

}

func (self nntpArticle) Subject() string {
  return self.headers.Get("Subject", "")
}

func (self nntpArticle) Posted() int64 {
  posted := self.headers.Get("Date", "")
  t, err := time.Parse(time.RFC1123Z, posted)
  if err == nil {
    return t.Unix()
  }
  return 0
}

func (self nntpArticle) Message() string {
  return self.message.body.String()
}

func (self nntpArticle) Path() string {
  return self.headers.Get("Path", "unspecified")
}

func (self nntpArticle) Headers() ArticleHeaders {
  return self.headers
}

func (self nntpArticle) AppendPath(part string) NNTPMessage {
  if self.headers.Has("Path") {
    self.headers.Set("Path", part + "!" + self.Path())
  } else {
    self.headers.Set("Path", part)
  }
  return self
}
func (self nntpArticle) ContentType() string {
  // assumes text/plain if unspecified
  return self.headers.Get("Content-Type", "text/plain")
}

func (self nntpArticle) Sage() bool {
  return self.headers.Has("X-Sage")
}

func (self nntpArticle) OP() bool {
  return self.headers.Get("Reference", self.headers.Get("References", "")) == ""
}

func (self nntpArticle) Attachments() []NNTPAttachment {
  return self.attachments
}

func (self nntpArticle) Attach(att NNTPAttachment) NNTPMessage {
  log.Println("attaching file", att.Filename())
  self.attachments = append(self.attachments, att)
  return self
}

func (self nntpArticle) WriteBody(wr io.Writer) (err error) {
  // this is a signed message, don't treat it special
  if self.signedPart.body.Len() > 0 {
    r := bufio.NewReader(&self.signedPart.body)
    w := bufio.NewWriter(wr)
    var line []byte
    for {
      // convert line endings :\
      line, err = r.ReadBytes('\n')
      if err == nil {
        w.Write(line[:len(line)-2])
        w.WriteByte(10)
      } else {
        return
      }
    }
  }
  if len(self.attachments) == 0 {
    // write plaintext and be done
    _, err = self.message.body.WriteTo(wr)
    return
  }
  content_type := self.ContentType()
  _, params, err := mime.ParseMediaType(content_type)
  if err != nil {
    log.Println("failed to parse media type", err)
    return err
  }

  boundary, ok := params["boundary"]
  if ok {
    w := multipart.NewWriter(wr)

    err = w.SetBoundary(boundary)
    if err == nil {
      attachments := []NNTPAttachment{self.message}
      attachments = append(attachments, self.attachments...)
      log.Println("writing nntp body")
      for _ , att := range(attachments) {
        hdr := att.Header()
        if hdr == nil {
          hdr = make(textproto.MIMEHeader)
        }
        hdr.Set("Content-Transfer-Encoding", "base64")
        part, err := w.CreatePart(hdr)
        enc := base64.NewEncoder(base64.StdEncoding, part)
        err = att.WriteTo(enc)
        enc.Close()

        if err != nil {
          break
        }
      }
    }
    if err != nil {
      log.Println(err)
    }
    err = w.Close()
  } else {
    _, err = self.message.body.WriteTo(wr)
  }
  return err
}
<|endoftext|>


<commit_before>package ecslogs

import "testing"

func TestParseFuncName(t *testing.T) {
	tests := []struct {
		name string
		pkg  string
		fn   string
	}{
		{
			name: "github.com/segmentio/ecs-logs/lib.TestLoggerPrintf",
			pkg:  "github.com/segmentio/ecs-logs/lib",
			fn:   "TestLoggerPrintf",
		},
		{
			name: "github.com/segmentio/ecs-logs/lib.(*Logger).Printf",
			pkg:  "github.com/segmentio/ecs-logs/lib",
			fn:   "(*Logger).Printf",
		},
		{
			name: "github.com/segmentio/ecs-logs/lib/apex.NewHandlerWith.func1",
			pkg:  "github.com/segmentio/ecs-logs/lib/apex",
			fn:   "NewHandlerWith.func1",
		},
		{
			name: "github.com/segmentio/ecs-logs/vendor/github.com/apex/log.HandlerFunc.HandleLog",
			pkg:  "github.com/segmentio/ecs-logs/vendor/github.com/apex/log",
			fn:   "HandlerFunc.HandleLog",
		},
	}

	for _, test := range tests {
		if pkg, fn := parseFuncName(test.name); pkg != test.pkg || fn != test.fn {
			t.Errorf("%s => %s - %s", test.name, pkg, fn)
		}
	}
}

func TestFuncInfoString(t *testing.T) {
	tests := []struct {
		info   FuncInfo
		source string
	}{
		{
			info:   FuncInfo{File: "test.go", Func: "F", Line: 42},
			source: "test.go:42:F",
		},
	}

	for _, test := range tests {
		if s := test.info.String(); s != test.source {
			t.Errorf("invalid source: %#v", s)
		}
	}
}
<commit_msg>add ecslogs.GuessCaller benchmark<commit_after>package ecslogs

import "testing"

func TestParseFuncName(t *testing.T) {
	tests := []struct {
		name string
		pkg  string
		fn   string
	}{
		{
			name: "github.com/segmentio/ecs-logs/lib.TestLoggerPrintf",
			pkg:  "github.com/segmentio/ecs-logs/lib",
			fn:   "TestLoggerPrintf",
		},
		{
			name: "github.com/segmentio/ecs-logs/lib.(*Logger).Printf",
			pkg:  "github.com/segmentio/ecs-logs/lib",
			fn:   "(*Logger).Printf",
		},
		{
			name: "github.com/segmentio/ecs-logs/lib/apex.NewHandlerWith.func1",
			pkg:  "github.com/segmentio/ecs-logs/lib/apex",
			fn:   "NewHandlerWith.func1",
		},
		{
			name: "github.com/segmentio/ecs-logs/vendor/github.com/apex/log.HandlerFunc.HandleLog",
			pkg:  "github.com/segmentio/ecs-logs/vendor/github.com/apex/log",
			fn:   "HandlerFunc.HandleLog",
		},
	}

	for _, test := range tests {
		if pkg, fn := parseFuncName(test.name); pkg != test.pkg || fn != test.fn {
			t.Errorf("%s => %s - %s", test.name, pkg, fn)
		}
	}
}

func TestFuncInfoString(t *testing.T) {
	tests := []struct {
		info   FuncInfo
		source string
	}{
		{
			info:   FuncInfo{File: "test.go", Func: "F", Line: 42},
			source: "test.go:42:F",
		},
	}

	for _, test := range tests {
		if s := test.info.String(); s != test.source {
			t.Errorf("invalid source: %#v", s)
		}
	}
}

func BenchmarkGuessCaller(b *testing.B) {
	for i := 0; i != b.N; i++ {
		GuessCaller(1, 10, "github.com/segmentio/ecs-logs")
	}
}
<|endoftext|>


<commit_before>// go-rst - A reStructuredText parser for Go
// 2014 (c) The go-rst Authors
// MIT Licensed. See LICENSE for details.

package parse

import (
	"fmt"
	"github.com/demizer/go-elog"
	"github.com/demizer/go-spew/spew"
	"os"
	"reflect"
)

var spd = spew.ConfigState{Indent: "\t", DisableMethods: true}

type systemMessageLevel int

const (
	levelInfo systemMessageLevel = iota
	levelWarning
	levelError
	levelSevere
)

var systemMessageLevels = [...]string{
	"INFO",
	"WARNING",
	"ERROR",
	"SEVERE",
}

func (s systemMessageLevel) String() string {
	return systemMessageLevels[s]
}

type systemMessage struct {
	level  systemMessageLevel
	line   int
	source string
	items  []item
}

type sectionLevels []*SectionNode

func (s *sectionLevels) String() string {
	var out string
	for _, sec := range *s {
		out += fmt.Sprintf("level: %d, rune: %q, overline: %t, length: %d\n",
			sec.Level, sec.UnderLine.Rune, sec.OverLine != nil, sec.Length)
	}
	return out
}

// Returns nil if not found
func (s *sectionLevels) FindByRune(adornChar rune) *SectionNode {
	for _, sec := range *s {
		if sec.UnderLine.Rune == adornChar {
			return sec
		}
	}
	return nil
}

// If exists == true, a section node with the same text and underline has been found in
// sectionLevels, sec is the matching SectionNode. If exists == false, then the sec return value is
// the similarly leveled SectionNode. If exists == false and sec == nil, then the SectionNode added
// to sectionLevels is a new Node.
func (s *sectionLevels) Add(section *SectionNode) (exists bool, sec *SectionNode) {
	sec = s.FindByRune(section.UnderLine.Rune)
	if sec != nil {
		if sec.Text == section.Text {
			return true, sec
		} else if sec.Text != section.Text {
			section.Level = sec.Level
		}
	} else {
		section.Level = len(*s) + 1
	}
	exists = false
	*s = append(*s, section)
	return
}

func (s *sectionLevels) Level() int {
	return len(*s)
}

// Parse is the entry point for the reStructuredText parser.
func Parse(name, text string) (t *Tree, errors []error) {
	t = New(name)
	t.text = text
	_, errors = t.Parse(text, t)
	return
}

func New(name string) *Tree {
	return &Tree{Name: name, Nodes: newList(), nodeTarget: newList(), sectionLevels: new(sectionLevels)}
}

var tokenPos = 2

type Tree struct {
	Name             string
	Nodes            *NodeList // The root node list
	nodeTarget       *NodeList // Used by the parser to add nodes to a target NodeList
	Errors           []error
	text             string
	lex              *lexer
	tokenBackupCount int
	tokenPeekCount   int
	token            [5]*item
	sectionLevels    *sectionLevels // Encountered section levels
	id               int            // The unique id of the node in the tree
}

func (t *Tree) errorf(format string, args ...interface{}) {
	format = fmt.Sprintf("go-rst: %s:%d: %s\n", t.Name, t.lex.lineNumber(), format)
	t.Errors = append(t.Errors, fmt.Errorf(format, args...))
}

func (t *Tree) error(err error) {
	t.errorf("%s\n", err)
}

// startParse initializes the parser, using the lexer.
func (t *Tree) startParse(lex *lexer) {
	t.lex = lex
}

// stopParse terminates parsing.
func (t *Tree) stopParse() {
	t.Nodes = nil
	t.nodeTarget = nil
	t.lex = nil
}

func (t *Tree) Parse(text string, treeSet *Tree) (tree *Tree, errors []error) {
	log.Debugln("Start")
	t.startParse(lex(t.Name, text))
	t.text = text
	t.parse(treeSet)
	log.Debugln("End")
	return t, t.Errors
}

func (t *Tree) parse(tree *Tree) {
	log.Debugln("Start")

	t.nodeTarget = t.Nodes

	for t.peek(1).Type != itemEOF {
		var n Node
		token := t.next()
		log.Infof("Got token: %#+v\n", token)

		switch token.Type {
		case itemTitle: // Section includes overline/underline
			n = t.section(token)
		case itemBlankLine:
			n = newBlankLine(token, &t.id)
		case itemParagraph:
			n = newParagraph(token, &t.id)
		case itemSpace:
			n = newSpace(token, &t.id)
		case itemSectionAdornment:
			// Section adornments should be consumed with itemTitle
			panic("Parser should not find itemSectionAdornment!")
		default:
			t.errorf("%q Not implemented!", token.Type)
			continue
		}

		t.nodeTarget.append(n)
		if n.NodeType() == NodeSection {
			t.nodeTarget =
				reflect.ValueOf(n).Elem().FieldByName("NodeList").Addr().Interface().(*NodeList)
		}
	}

	log.Debugln("End")
}

func (t *Tree) backup() *item {
	t.tokenBackupCount++
	// log.Debugln("t.tokenBackupCount:", t.tokenPeekCount)
	if t.tokenBackupCount > 2 {
		panic("t.backup() can only be used twice consecutively.")
	}
	for i := 4; i > 0; i-- {
		t.token[i] = t.token[i-1]
		t.token[i-1] = nil
	}
	// log.Debugf("\n##### backup() aftermath #####\n\n")
	// spd.Dump(t.token)
	return t.token[tokenPos-t.tokenBackupCount]
}

func (t *Tree) peekBack(pos int) *item {
	if pos > 2 {
		panic("Cannot peek back more than two positions!")
	}
	return t.token[tokenPos-pos]
}

func (t *Tree) peek(pos int) *item {
	if pos > 2 {
		panic("It is only possible to peek ahead two positions!")
	}
	for i := pos; i <= pos; i++ {
		log.Debugln(i)
		t.tokenPeekCount++
		if t.token[tokenPos+t.tokenPeekCount] == nil {
			t.token[tokenPos+t.tokenPeekCount] = t.lex.nextItem()
		}
	}
	return t.token[tokenPos+t.tokenPeekCount]
}

func (t *Tree) next() *item {
	// shifts the pointers left in t.token, pos is the amount to shift
	shift := func(pos int) {
		for i := pos; i > 0; i-- {
			if t.tokenPeekCount > 0 {
				if t.token[tokenPos+t.tokenPeekCount+1] != nil {
					panic("t.token[t.tokenPeekCount] should be nil!")
				}
				t.token[tokenPos+t.tokenPeekCount+1] = t.lex.nextItem()
				t.tokenPeekCount--
			}
			for x := 0; x < 4; x++ {
				t.token[x] = t.token[x+1]
				t.token[x+1] = nil
			}
		}
	}
	if t.tokenPeekCount > 0 {
		shift(t.tokenPeekCount)
	} else {
		shift(1)
	}
	t.tokenBackupCount, t.tokenPeekCount = 0, 0
	return t.token[tokenPos]
}

func (t *Tree) section(i *item) Node {
	log.Debugln("Start")
	var overAdorn, title, underAdorn *item
	var overline bool

	peekBack := t.peekBack(1)
	if peekBack != nil {
		switch peekBack.Type {
		case itemSectionAdornment:
			overline = true
			overAdorn = peekBack
		case itemSpace:
			// TODO: Handle indented titles here!
			log.Debugln("FOUND ITEMSPACE BEFORE TITLE")
		}
	}

	title = i
	underAdorn = t.next() // Grab the section underline

	// Check adornment for proper syntax
	if title.Length != underAdorn.Length {
		t.errorf("Section under line  not equal to title length!")
	} else if overline && title.Length != overAdorn.Length {
		t.errorf("Section over line not equal to title length!")
	} else if overline && overAdorn.Text != underAdorn.Text {
		t.errorf("Section title over line does not match section title under line.")
	}

	sec := newSection(title, &t.id, overAdorn, underAdorn)
	exists, eSec := t.sectionLevels.Add(sec)
	if exists && eSec != nil {
		t.errorf("SectionNode using Text \"%s\" and Rune '%s' was previously parsed!",
			sec.Text, string(sec.UnderLine.Rune))
	} else if !exists && eSec != nil {
		// There is a matching level in sectionLevels
		t.nodeTarget = &(*t.sectionLevels)[sec.Level - 2].NodeList
	}

	log.Debugln("End")
	return sec
}
<commit_msg>parse.go: Adjust logging in t.peek()<commit_after>// go-rst - A reStructuredText parser for Go
// 2014 (c) The go-rst Authors
// MIT Licensed. See LICENSE for details.

package parse

import (
	"fmt"
	"github.com/demizer/go-elog"
	"github.com/demizer/go-spew/spew"
	"os"
	"reflect"
)

var spd = spew.ConfigState{Indent: "\t", DisableMethods: true}

type systemMessageLevel int

const (
	levelInfo systemMessageLevel = iota
	levelWarning
	levelError
	levelSevere
)

var systemMessageLevels = [...]string{
	"INFO",
	"WARNING",
	"ERROR",
	"SEVERE",
}

func (s systemMessageLevel) String() string {
	return systemMessageLevels[s]
}

type systemMessage struct {
	level  systemMessageLevel
	line   int
	source string
	items  []item
}

type sectionLevels []*SectionNode

func (s *sectionLevels) String() string {
	var out string
	for _, sec := range *s {
		out += fmt.Sprintf("level: %d, rune: %q, overline: %t, length: %d\n",
			sec.Level, sec.UnderLine.Rune, sec.OverLine != nil, sec.Length)
	}
	return out
}

// Returns nil if not found
func (s *sectionLevels) FindByRune(adornChar rune) *SectionNode {
	for _, sec := range *s {
		if sec.UnderLine.Rune == adornChar {
			return sec
		}
	}
	return nil
}

// If exists == true, a section node with the same text and underline has been found in
// sectionLevels, sec is the matching SectionNode. If exists == false, then the sec return value is
// the similarly leveled SectionNode. If exists == false and sec == nil, then the SectionNode added
// to sectionLevels is a new Node.
func (s *sectionLevels) Add(section *SectionNode) (exists bool, sec *SectionNode) {
	sec = s.FindByRune(section.UnderLine.Rune)
	if sec != nil {
		if sec.Text == section.Text {
			return true, sec
		} else if sec.Text != section.Text {
			section.Level = sec.Level
		}
	} else {
		section.Level = len(*s) + 1
	}
	exists = false
	*s = append(*s, section)
	return
}

func (s *sectionLevels) Level() int {
	return len(*s)
}

// Parse is the entry point for the reStructuredText parser.
func Parse(name, text string) (t *Tree, errors []error) {
	t = New(name)
	t.text = text
	_, errors = t.Parse(text, t)
	return
}

func New(name string) *Tree {
	return &Tree{Name: name, Nodes: newList(), nodeTarget: newList(), sectionLevels: new(sectionLevels)}
}

var tokenPos = 2

type Tree struct {
	Name             string
	Nodes            *NodeList // The root node list
	nodeTarget       *NodeList // Used by the parser to add nodes to a target NodeList
	Errors           []error
	text             string
	lex              *lexer
	tokenBackupCount int
	tokenPeekCount   int
	token            [5]*item
	sectionLevels    *sectionLevels // Encountered section levels
	id               int            // The unique id of the node in the tree
}

func (t *Tree) errorf(format string, args ...interface{}) {
	format = fmt.Sprintf("go-rst: %s:%d: %s\n", t.Name, t.lex.lineNumber(), format)
	t.Errors = append(t.Errors, fmt.Errorf(format, args...))
}

func (t *Tree) error(err error) {
	t.errorf("%s\n", err)
}

// startParse initializes the parser, using the lexer.
func (t *Tree) startParse(lex *lexer) {
	t.lex = lex
}

// stopParse terminates parsing.
func (t *Tree) stopParse() {
	t.Nodes = nil
	t.nodeTarget = nil
	t.lex = nil
}

func (t *Tree) Parse(text string, treeSet *Tree) (tree *Tree, errors []error) {
	log.Debugln("Start")
	t.startParse(lex(t.Name, text))
	t.text = text
	t.parse(treeSet)
	log.Debugln("End")
	return t, t.Errors
}

func (t *Tree) parse(tree *Tree) {
	log.Debugln("Start")

	t.nodeTarget = t.Nodes

	for t.peek(1).Type != itemEOF {
		var n Node
		token := t.next()
		log.Infof("Got token: %#+v\n", token)

		switch token.Type {
		case itemTitle: // Section includes overline/underline
			n = t.section(token)
		case itemBlankLine:
			n = newBlankLine(token, &t.id)
		case itemParagraph:
			n = newParagraph(token, &t.id)
		case itemSpace:
			n = newSpace(token, &t.id)
		case itemSectionAdornment:
			// Section adornments should be consumed with itemTitle
			panic("Parser should not find itemSectionAdornment!")
		default:
			t.errorf("%q Not implemented!", token.Type)
			continue
		}

		t.nodeTarget.append(n)
		if n.NodeType() == NodeSection {
			t.nodeTarget =
				reflect.ValueOf(n).Elem().FieldByName("NodeList").Addr().Interface().(*NodeList)
		}
	}

	log.Debugln("End")
}

func (t *Tree) backup() *item {
	t.tokenBackupCount++
	// log.Debugln("t.tokenBackupCount:", t.tokenPeekCount)
	if t.tokenBackupCount > 2 {
		panic("t.backup() can only be used twice consecutively.")
	}
	for i := 4; i > 0; i-- {
		t.token[i] = t.token[i-1]
		t.token[i-1] = nil
	}
	// log.Debugf("\n##### backup() aftermath #####\n\n")
	// spd.Dump(t.token)
	return t.token[tokenPos-t.tokenBackupCount]
}

func (t *Tree) peekBack(pos int) *item {
	if pos > 2 {
		panic("Cannot peek back more than two positions!")
	}
	return t.token[tokenPos-pos]
}

func (t *Tree) peek(pos int) *item {
	// log.Debugln("t.tokenPeekCount:", t.tokenPeekCount)
	if pos > 2 {
		panic("It is only possible to peek ahead two positions!")
	}
	for i := pos; i <= pos; i++ {
		t.tokenPeekCount++
		if t.token[tokenPos+t.tokenPeekCount] == nil {
			t.token[tokenPos+t.tokenPeekCount] = t.lex.nextItem()
		}
	}
	// log.Debugf("\n##### peek() aftermath #####\n\n")
	// spd.Dump(t.token)
	return t.token[tokenPos+t.tokenPeekCount]
}

func (t *Tree) next() *item {
	// shifts the pointers left in t.token, pos is the amount to shift
	shift := func(pos int) {
		for i := pos; i > 0; i-- {
			if t.tokenPeekCount > 0 {
				if t.token[tokenPos+t.tokenPeekCount+1] != nil {
					panic("t.token[t.tokenPeekCount] should be nil!")
				}
				t.token[tokenPos+t.tokenPeekCount+1] = t.lex.nextItem()
				t.tokenPeekCount--
			}
			for x := 0; x < 4; x++ {
				t.token[x] = t.token[x+1]
				t.token[x+1] = nil
			}
		}
	}
	if t.tokenPeekCount > 0 {
		shift(t.tokenPeekCount)
	} else {
		shift(1)
	}
	t.tokenBackupCount, t.tokenPeekCount = 0, 0
	return t.token[tokenPos]
}

func (t *Tree) section(i *item) Node {
	log.Debugln("Start")
	var overAdorn, title, underAdorn *item
	var overline bool

	peekBack := t.peekBack(1)
	if peekBack != nil {
		switch peekBack.Type {
		case itemSectionAdornment:
			overline = true
			overAdorn = peekBack
		case itemSpace:
			// TODO: Handle indented titles here!
			log.Debugln("FOUND ITEMSPACE BEFORE TITLE")
		}
	}

	title = i
	underAdorn = t.next() // Grab the section underline

	// Check adornment for proper syntax
	if title.Length != underAdorn.Length {
		t.errorf("Section under line  not equal to title length!")
	} else if overline && title.Length != overAdorn.Length {
		t.errorf("Section over line not equal to title length!")
	} else if overline && overAdorn.Text != underAdorn.Text {
		t.errorf("Section title over line does not match section title under line.")
	}

	sec := newSection(title, &t.id, overAdorn, underAdorn)
	exists, eSec := t.sectionLevels.Add(sec)
	if exists && eSec != nil {
		t.errorf("SectionNode using Text \"%s\" and Rune '%s' was previously parsed!",
			sec.Text, string(sec.UnderLine.Rune))
	} else if !exists && eSec != nil {
		// There is a matching level in sectionLevels
		t.nodeTarget = &(*t.sectionLevels)[sec.Level - 2].NodeList
	}

	log.Debugln("End")
	return sec
}
<|endoftext|>


<commit_before>package parser

import (
	"errors"

	"github.com/Sirupsen/logrus"
	"github.com/onsi/gocleanup"
)

/*
Lint parses a builderfile and returns either nil if the file was parsed
successfully or an error indicating that parsing failed and the file is
invalid.
*/
func (parser *Parser) Lint() error {
	//TODO: this should probably call Parse instead of the unexported method rawToStruct
	_, err := parser.rawToStruct()

	if err != nil {
		return err
	}

	return nil
}

/*
AssertLint is like Lint except that instead of returning an nil/error to
indicate success/failure, it exits nonzero if linting fails.
*/
func (parser *Parser) AssertLint() {
	if !parser.IsOpenable() {
		if parser.filename == "" {
			parser.printLintFailMessage(errors.New("no file provided for linting"))
		} else {
			parser.printLintFailMessage(errors.New("unable to open file"))
		}
		gocleanup.Exit(17)
	}

	err := parser.Lint()
	if err != nil {
		parser.printLintFailMessage(err)
		gocleanup.Exit(5)
	} else {
		parser.printLintSuccessMessage()
		gocleanup.Exit(0)
	}
}

// helper functions
func (parser *Parser) printLintSuccessMessage() {
	//parser.WithFields(logrus.Fields{"file": parser.filename}).Info("file is a valid Bobfile")
}

func (parser *Parser) printLintFailMessage(err error) {
	//var errMsg string
	fields := logrus.Fields{"error": err}

	if parser.filename == "" {
		//errMsg = "no file provided for linting"
	} else {
		fields["filename"] = parser.filename
		//errMsg = "file provided is not a valid Bobfile"
	}
	//parser.WithFields(fields).Error(errMsg)
}
<commit_msg>Forget why I commented out logging for lint functionality<commit_after>package parser

import (
	"errors"

	"github.com/Sirupsen/logrus"
	"github.com/onsi/gocleanup"
)

/*
Lint parses a builderfile and returns either nil if the file was parsed
successfully or an error indicating that parsing failed and the file is
invalid.
*/
func (parser *Parser) Lint() error {
	//TODO: this should probably call Parse instead of the unexported method rawToStruct
	_, err := parser.rawToStruct()
	return err
}

/*
AssertLint is like Lint except that instead of returning an nil/error to
indicate success/failure, it exits nonzero if linting fails.
*/
func (parser *Parser) AssertLint() {
	if !parser.IsOpenable() {
		if parser.filename == "" {
			parser.printLintFailMessage(errors.New("no file provided for linting"))
		} else {
			parser.printLintFailMessage(errors.New("unable to open file"))
		}
		gocleanup.Exit(17)
	}

	err := parser.Lint()
	if err != nil {
		parser.printLintFailMessage(err)
		gocleanup.Exit(5)
	} else {
		parser.printLintSuccessMessage()
		gocleanup.Exit(0)
	}
}

// helper functions
func (parser *Parser) printLintSuccessMessage() {
	parser.WithFields(logrus.Fields{"file": parser.filename}).Info("file is a valid Bobfile")
}

func (parser *Parser) printLintFailMessage(err error) {
	var errMsg string
	fields := logrus.Fields{"error": err}

	if parser.filename == "" {
		errMsg = "no file provided for linting"
	} else {
		fields["filename"] = parser.filename
		errMsg = "file provided is not a valid Bobfile"
	}
	parser.WithFields(fields).Error(errMsg)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"gopkg.in/h2non/gentleman.v1"
	"gopkg.in/h2non/gentleman.v1/context"
)

// Usege example with gentleman HTTP client toolkit.
// See also: https://github.com/h2non/gentleman-mock
func main() {
	defer gock.Off()

	gock.New("http://httpbin.org").
		Get("/*").
		Reply(204).
		SetHeader("Server", "gock")

	cli := gentleman.New()

	cli.UseHandler("before dial", func(ctx *context.Context, h context.Handler) {
		gock.InterceptClient(ctx.Client)
		h.Next(ctx)
	})

	res, err := cli.Request().URL("http://httpbin.org/get").Send()
	if err != nil {
		fmt.Errorf("Error: %s", err)
	}

	fmt.Printf("Status: %d\n", res.StatusCode)
	fmt.Printf("Server header: %s\n", res.Header.Get("Server"))
}
<commit_msg>fix(examples): add imports<commit_after>package main

import (
	"fmt"

	"gopkg.in/h2non/gock.v1"
	"gopkg.in/h2non/gentleman.v1"
	"gopkg.in/h2non/gentleman.v1/context"
)

// Usege example with gentleman HTTP client toolkit.
// See also: https://github.com/h2non/gentleman-mock
func main() {
	defer gock.Off()

	gock.New("http://httpbin.org").
		Get("/*").
		Reply(204).
		SetHeader("Server", "gock")

	cli := gentleman.New()

	cli.UseHandler("before dial", func(ctx *context.Context, h context.Handler) {
		gock.InterceptClient(ctx.Client)
		h.Next(ctx)
	})

	res, err := cli.Request().URL("http://httpbin.org/get").Send()
	if err != nil {
		fmt.Errorf("Error: %s", err)
	}

	fmt.Printf("Status: %d\n", res.StatusCode)
	fmt.Printf("Server header: %s\n", res.Header.Get("Server"))
}
<|endoftext|>


<commit_before>package gedis

import (
	"testing"
	"bytes"
	"strings"
	"runtime"
)

func assertParse(t *testing.T, input, output string) {
	_, file, ln, _ := runtime.Caller(1)

	reader := strings.NewReader(input)
	expected := []byte(output)

	data, err := Parse(reader)

	if err != nil {
		t.Errorf("%s:%d: Parse(%#v) returned an error: %v", file, ln, []byte(input), err)
		t.FailNow()
	}

	got, ok := data.([]byte)

	if !ok {
		t.Errorf("%s:%d: Cannot convert to []byte: %#v", file, ln, data)
		t.FailNow()
	}

	if !bytes.Equal(expected, got) {
		t.Errorf("%s:%d: Parse(%#v)\nreturned %#v\nexpected %#v", file, ln, []byte(input), data, expected)
	}
}

func Test_Parse_Status(t *testing.T) {
	assertParse(t, "+OK\r\n", "OK")
	assertParse(t, "+PING\r\n", "PING")
}

func Benchmark_Parse_Status(b *testing.B) {
	for i := 0; i < b.N; i++ {
		b.StopTimer()
		reader := strings.NewReader("+OK\r\n")
		b.StartTimer()
		Parse(reader)
	}
}

func Test_Parse_Error(t *testing.T) {
	_, err := Parse(strings.NewReader("-ERR unknown\r\n"))

	if err == nil {
		t.Errorf("Parsing an error reply didn't return an error")
		t.FailNow()
	}

	if err.Error() != "ERR unknown" {
		t.Errorf("Unexpected error: %v", err)
		t.FailNow()
	}
}

func assertInteger(t *testing.T, input string, expected int) {
	_, file, ln, _ := runtime.Caller(1)

	reader := strings.NewReader(input)

	ret, err := Parse(reader)

	if err != nil {
		t.Errorf("%s:%d: Parse(%#v) returned an error: %v", file, ln, []byte(input), err)
		t.FailNow()
	}

	got, ok := ret.(int)

	if !ok {
		t.Errorf("%s:%d: Parse(%q): Couldn't convert to int: %#v", file, ln, input, ret)
		t.FailNow()
	}

	if got != expected {
		t.Errorf("%s:%d: Parse(%#v)\nreturned %#v\nexpected %#v", file, ln, []byte(input), got, expected)
		t.FailNow()
	}
}

func Test_Parse_Integer(t *testing.T) {
	assertInteger(t, ":1234\r\n", 1234)
	assertInteger(t, ":-1234\r\n", -1234)
}

func Test_Parse_Bulk(t *testing.T) {
	assertParse(t, "$5\r\nlorem\r\n", "lorem")
	assertParse(t, "$12\r\nlorem\r\nipsum\r\n", "lorem\r\nipsum")
}
<commit_msg>Rename helper function<commit_after>package gedis

import (
	"testing"
	"bytes"
	"strings"
	"runtime"
)

func assertString(t *testing.T, input, output string) {
	_, file, ln, _ := runtime.Caller(1)

	reader := strings.NewReader(input)
	expected := []byte(output)

	data, err := Parse(reader)

	if err != nil {
		t.Errorf("%s:%d: Parse(%#v) returned an error: %v", file, ln, []byte(input), err)
		t.FailNow()
	}

	got, ok := data.([]byte)

	if !ok {
		t.Errorf("%s:%d: Cannot convert to []byte: %#v", file, ln, data)
		t.FailNow()
	}

	if !bytes.Equal(expected, got) {
		t.Errorf("%s:%d: Parse(%#v)\nreturned %#v\nexpected %#v", file, ln, []byte(input), data, expected)
	}
}

func Test_Parse_Status(t *testing.T) {
	assertString(t, "+OK\r\n", "OK")
	assertString(t, "+PING\r\n", "PING")
}

func Benchmark_Parse_Status(b *testing.B) {
	for i := 0; i < b.N; i++ {
		b.StopTimer()
		reader := strings.NewReader("+OK\r\n")
		b.StartTimer()
		Parse(reader)
	}
}

func Test_Parse_Error(t *testing.T) {
	_, err := Parse(strings.NewReader("-ERR unknown\r\n"))

	if err == nil {
		t.Errorf("Parsing an error reply didn't return an error")
		t.FailNow()
	}

	if err.Error() != "ERR unknown" {
		t.Errorf("Unexpected error: %v", err)
		t.FailNow()
	}
}

func assertInteger(t *testing.T, input string, expected int) {
	_, file, ln, _ := runtime.Caller(1)

	reader := strings.NewReader(input)

	ret, err := Parse(reader)

	if err != nil {
		t.Errorf("%s:%d: Parse(%#v) returned an error: %v", file, ln, []byte(input), err)
		t.FailNow()
	}

	got, ok := ret.(int)

	if !ok {
		t.Errorf("%s:%d: Parse(%q): Couldn't convert to int: %#v", file, ln, input, ret)
		t.FailNow()
	}

	if got != expected {
		t.Errorf("%s:%d: Parse(%#v)\nreturned %#v\nexpected %#v", file, ln, []byte(input), got, expected)
		t.FailNow()
	}
}

func Test_Parse_Integer(t *testing.T) {
	assertInteger(t, ":1234\r\n", 1234)
	assertInteger(t, ":-1234\r\n", -1234)
}

func Test_Parse_Bulk(t *testing.T) {
	assertString(t, "$5\r\nlorem\r\n", "lorem")
	assertString(t, "$12\r\nlorem\r\nipsum\r\n", "lorem\r\nipsum")
}
<|endoftext|>


<commit_before>package nats_output

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

	"github.com/google/uuid"
	"github.com/karimra/gnmic/outputs"
	"github.com/mitchellh/mapstructure"
	"github.com/nats-io/nats.go"
	"github.com/prometheus/client_golang/prometheus"
)

const (
	natsConnectTimeout = 5 * time.Second
	natsConnectWait    = 2 * time.Second

	natsReconnectBufferSize = 100 * 1024 * 1024
	natsReconnectWait       = 2 * time.Second

	defaultSubjectName = "gnmic-telemetry"
)

func init() {
	outputs.Register("nats", func() outputs.Output {
		return &NatsOutput{
			Cfg: &Config{},
		}
	})
}

// NatsOutput //
type NatsOutput struct {
	Cfg      *Config
	ctx      context.Context
	cancelFn context.CancelFunc
	conn     *nats.Conn
	metrics  []prometheus.Collector
	logger   *log.Logger
}

// Config //
type Config struct {
	Name            string        `mapstructure:"name,omitempty"`
	Address         string        `mapstructure:"address,omitempty"`
	SubjectPrefix   string        `mapstructure:"subject-prefix,omitempty"`
	Subject         string        `mapstructure:"subject,omitempty"`
	Username        string        `mapstructure:"username,omitempty"`
	Password        string        `mapstructure:"password,omitempty"`
	ConnectTimeWait time.Duration `mapstructure:"connect-time-wait,omitempty"`
}

func (n *NatsOutput) String() string {
	b, err := json.Marshal(n)
	if err != nil {
		return ""
	}
	return string(b)
}

// Init //
func (n *NatsOutput) Init(cfg map[string]interface{}, logger *log.Logger) error {
	err := mapstructure.Decode(cfg, n.Cfg)
	if err != nil {
		return err
	}
	if n.Cfg.ConnectTimeWait == 0 {
		n.Cfg.ConnectTimeWait = natsConnectWait
	}
	if n.Cfg.Subject == "" && n.Cfg.SubjectPrefix == "" {
		n.Cfg.Subject = defaultSubjectName
	}
	n.logger = log.New(os.Stderr, "nats_output ", log.LstdFlags|log.Lmicroseconds)
	if logger != nil {
		n.logger.SetOutput(logger.Writer())
		n.logger.SetFlags(logger.Flags())
	}
	if n.Cfg.Name == "" {
		n.Cfg.Name = "gnmic-" + uuid.New().String()
	}
	n.ctx, n.cancelFn = context.WithCancel(context.Background())
	n.conn, err = n.createNATSConn(n.Cfg)
	if err != nil {
		return err
	}
	n.logger.Printf("initialized nats producer: %s", n.String())
	return nil
}

// Write //
func (n *NatsOutput) Write(b []byte, meta outputs.Meta) {
	if len(b) == 0 {
		return
	}
	if format, ok := meta["format"]; ok {
		if format == "textproto" {
			return
		}
	}
	ssb := strings.Builder{}
	ssb.WriteString(n.Cfg.SubjectPrefix)
	if n.Cfg.SubjectPrefix != "" {
		if s, ok := meta["source"]; ok {
			source := strings.ReplaceAll(fmt.Sprintf("%s", s), ".", "-")
			source = strings.ReplaceAll(source, " ", "_")
			ssb.WriteString(".")
			ssb.WriteString(source)
		}
		if subname, ok := meta["subscription-name"]; ok {
			ssb.WriteString(".")
			ssb.WriteString(fmt.Sprintf("%s", subname))
		}
	} else if n.Cfg.Subject != "" {
		ssb.WriteString(n.Cfg.Subject)
	}
	subject := strings.ReplaceAll(ssb.String(), " ", "_")
	err := n.conn.Publish(subject, b)
	if err != nil {
		log.Printf("failed to write to nats subject '%s': %v", subject, err)
		return
	}
	// n.logger.Printf("wrote %d bytes to nats_subject=%s", len(b), n.Cfg.Subject)
}

// Close //
func (n *NatsOutput) Close() error {
	n.cancelFn()
	n.conn.Close()
	return nil
}

// Metrics //
func (n *NatsOutput) Metrics() []prometheus.Collector { return n.metrics }

func (n *NatsOutput) createNATSConn(c *Config) (*nats.Conn, error) {
	opts := []nats.Option{
		nats.Name(c.Name),
		nats.SetCustomDialer(n),
		nats.ReconnectWait(natsReconnectWait),
		nats.ReconnectBufSize(natsReconnectBufferSize),
		nats.ErrorHandler(func(_ *nats.Conn, _ *nats.Subscription, err error) {
			n.logger.Printf("NATS error: %v", err)
		}),
		nats.DisconnectHandler(func(c *nats.Conn) {
			n.logger.Println("Disconnected from NATS")
		}),
		nats.ClosedHandler(func(c *nats.Conn) {
			n.logger.Println("NATS connection is closed")
		}),
	}
	if c.Username != "" && c.Password != "" {
		opts = append(opts, nats.UserInfo(c.Username, c.Password))
	}
	nc, err := nats.Connect(c.Address, opts...)
	if err != nil {
		return nil, err
	}
	return nc, nil
}

// Dial //
func (n *NatsOutput) Dial(network, address string) (net.Conn, error) {
	ctx, cancel := context.WithCancel(n.ctx)
	defer cancel()

	for {
		n.logger.Println("attempting to connect to", address)
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}

		select {
		case <-n.ctx.Done():
			return nil, n.ctx.Err()
		default:
			d := &net.Dialer{}
			if conn, err := d.DialContext(ctx, network, address); err == nil {
				n.logger.Printf("successfully connected to NATS server %s", address)
				return conn, nil
			}
			time.Sleep(n.Cfg.ConnectTimeWait)
		}
	}
}
<commit_msg>use configured connect time wait<commit_after>package nats_output

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

	"github.com/google/uuid"
	"github.com/karimra/gnmic/outputs"
	"github.com/mitchellh/mapstructure"
	"github.com/nats-io/nats.go"
	"github.com/prometheus/client_golang/prometheus"
)

const (
	natsConnectTimeout = 5 * time.Second
	natsConnectWait    = 2 * time.Second

	natsReconnectBufferSize = 100 * 1024 * 1024
	natsReconnectWait       = 2 * time.Second

	defaultSubjectName = "gnmic-telemetry"
)

func init() {
	outputs.Register("nats", func() outputs.Output {
		return &NatsOutput{
			Cfg: &Config{},
		}
	})
}

// NatsOutput //
type NatsOutput struct {
	Cfg      *Config
	ctx      context.Context
	cancelFn context.CancelFunc
	conn     *nats.Conn
	metrics  []prometheus.Collector
	logger   *log.Logger
}

// Config //
type Config struct {
	Name            string        `mapstructure:"name,omitempty"`
	Address         string        `mapstructure:"address,omitempty"`
	SubjectPrefix   string        `mapstructure:"subject-prefix,omitempty"`
	Subject         string        `mapstructure:"subject,omitempty"`
	Username        string        `mapstructure:"username,omitempty"`
	Password        string        `mapstructure:"password,omitempty"`
	ConnectTimeWait time.Duration `mapstructure:"connect-time-wait,omitempty"`
}

func (n *NatsOutput) String() string {
	b, err := json.Marshal(n)
	if err != nil {
		return ""
	}
	return string(b)
}

// Init //
func (n *NatsOutput) Init(cfg map[string]interface{}, logger *log.Logger) error {
	err := mapstructure.Decode(cfg, n.Cfg)
	if err != nil {
		return err
	}
	if n.Cfg.ConnectTimeWait == 0 {
		n.Cfg.ConnectTimeWait = natsConnectWait
	}
	if n.Cfg.Subject == "" && n.Cfg.SubjectPrefix == "" {
		n.Cfg.Subject = defaultSubjectName
	}
	n.logger = log.New(os.Stderr, "nats_output ", log.LstdFlags|log.Lmicroseconds)
	if logger != nil {
		n.logger.SetOutput(logger.Writer())
		n.logger.SetFlags(logger.Flags())
	}
	if n.Cfg.Name == "" {
		n.Cfg.Name = "gnmic-" + uuid.New().String()
	}
	n.ctx, n.cancelFn = context.WithCancel(context.Background())
	n.conn, err = n.createNATSConn(n.Cfg)
	if err != nil {
		return err
	}
	n.logger.Printf("initialized nats producer: %s", n.String())
	return nil
}

// Write //
func (n *NatsOutput) Write(b []byte, meta outputs.Meta) {
	if len(b) == 0 {
		return
	}
	if format, ok := meta["format"]; ok {
		if format == "textproto" {
			return
		}
	}
	ssb := strings.Builder{}
	ssb.WriteString(n.Cfg.SubjectPrefix)
	if n.Cfg.SubjectPrefix != "" {
		if s, ok := meta["source"]; ok {
			source := strings.ReplaceAll(fmt.Sprintf("%s", s), ".", "-")
			source = strings.ReplaceAll(source, " ", "_")
			ssb.WriteString(".")
			ssb.WriteString(source)
		}
		if subname, ok := meta["subscription-name"]; ok {
			ssb.WriteString(".")
			ssb.WriteString(fmt.Sprintf("%s", subname))
		}
	} else if n.Cfg.Subject != "" {
		ssb.WriteString(n.Cfg.Subject)
	}
	subject := strings.ReplaceAll(ssb.String(), " ", "_")
	err := n.conn.Publish(subject, b)
	if err != nil {
		log.Printf("failed to write to nats subject '%s': %v", subject, err)
		return
	}
	// n.logger.Printf("wrote %d bytes to nats_subject=%s", len(b), n.Cfg.Subject)
}

// Close //
func (n *NatsOutput) Close() error {
	n.cancelFn()
	n.conn.Close()
	return nil
}

// Metrics //
func (n *NatsOutput) Metrics() []prometheus.Collector { return n.metrics }

func (n *NatsOutput) createNATSConn(c *Config) (*nats.Conn, error) {
	opts := []nats.Option{
		nats.Name(c.Name),
		nats.SetCustomDialer(n),
		nats.ReconnectWait(n.Cfg.ConnectTimeWait),
		nats.ReconnectBufSize(natsReconnectBufferSize),
		nats.ErrorHandler(func(_ *nats.Conn, _ *nats.Subscription, err error) {
			n.logger.Printf("NATS error: %v", err)
		}),
		nats.DisconnectHandler(func(c *nats.Conn) {
			n.logger.Println("Disconnected from NATS")
		}),
		nats.ClosedHandler(func(c *nats.Conn) {
			n.logger.Println("NATS connection is closed")
		}),
	}
	if c.Username != "" && c.Password != "" {
		opts = append(opts, nats.UserInfo(c.Username, c.Password))
	}
	nc, err := nats.Connect(c.Address, opts...)
	if err != nil {
		return nil, err
	}
	return nc, nil
}

// Dial //
func (n *NatsOutput) Dial(network, address string) (net.Conn, error) {
	ctx, cancel := context.WithCancel(n.ctx)
	defer cancel()

	for {
		n.logger.Println("attempting to connect to", address)
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}

		select {
		case <-n.ctx.Done():
			return nil, n.ctx.Err()
		default:
			d := &net.Dialer{}
			if conn, err := d.DialContext(ctx, network, address); err == nil {
				n.logger.Printf("successfully connected to NATS server %s", address)
				return conn, nil
			}
			time.Sleep(n.Cfg.ConnectTimeWait)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// A program that makes random reads within a large object on GCS, reporting
// throughput and latency information.
package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"os/signal"
	"sync"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"

	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/syncutil"
)

var fBucket = flag.String("bucket", "", "The GCS bucket from which to read.")
var fObject = flag.String("object", "", "The object within which to read.")
var fSize = flag.Uint64("size", 1<<21, "The size of each read in bytes.")
var fWorkers = flag.Int("workers", 1, "The number of workers to run.")

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type result struct {
	BytesRead int

	// Time taken from making the HTTP request to receiving the first byte of the
	// response body.
	FirstByteLatency time.Duration

	// Time taken from making the HTTP request to receiving the last byte of the
	// response body.
	FullBodyDuration time.Duration
}

func readOnce(
	ctx context.Context,
	o *gcs.Object,
	bucket gcs.Bucket) (r result, err error) {
	// Is the object large enough?
	if o.Size < *fSize {
		err = fmt.Errorf(
			"Object of size %d not large enough for read size %d",
			o.Size,
			*fSize)

		return
	}

	// Set up an appropriate request.
	req := &gcs.ReadObjectRequest{
		Name:       o.Name,
		Generation: o.Generation,
		Range:      &gcs.ByteRange{},
	}

	req.Range.Start = uint64(rand.Int63n(int64(o.Size - *fSize)))
	req.Range.Limit = req.Range.Start + *fSize

	// Create the reader.
	start := time.Now()
	rc, err := bucket.NewReader(ctx, req)
	if err != nil {
		err = fmt.Errorf("NewReader: %v", err)
		return
	}

	defer func() {
		closeErr := rc.Close()
		if err == nil && closeErr != nil {
			err = fmt.Errorf("Close: %v", closeErr)
		}
	}()

	// Measure the time to first byte.
	_, err = rc.Read([]byte{0})
	if err != nil {
		err = fmt.Errorf("Read: %v", err)
		return
	}

	r.FirstByteLatency = time.Since(start)

	// And the time to read everything.
	n, err := io.Copy(ioutil.Discard, rc)
	if err != nil {
		err = fmt.Errorf("Copy: %v", err)
		return
	}

	r.FullBodyDuration = time.Since(start)
	r.BytesRead = int(n + 1)

	return
}

// Make random reads within the given object, writing results to the supplied
// channel. Stop when the stop channel is closed.
func makeReads(
	ctx context.Context,
	o *gcs.Object,
	bucket gcs.Bucket,
	results chan<- result,
	stop <-chan struct{}) (err error) {
	// Is the object large enough?
	if o.Size < *fSize {
		err = fmt.Errorf(
			"Object of size %d not large enough for read size %d",
			o.Size,
			*fSize)

		return
	}

	for {
		// Stop?
		if _, ok := <-stop; ok {
			return
		}

		// Make a read.
		var r result
		r, err = readOnce(ctx, o, bucket)
		if err != nil {
			err = fmt.Errorf("readOnce %v", err)
			return
		}

		log.Printf(
			"Read %d bytes in %v (first byte latency %v)",
			r.BytesRead,
			r.FullBodyDuration,
			r.FirstByteLatency)

		// Write out the result.
		select {
		case <-ctx.Done():
			err = ctx.Err()
			return

		case results <- r:
		}
	}
}

func getBucket(ctx context.Context) (b gcs.Bucket, err error) {
	if *fBucket == "" {
		err = errors.New("You must set --bucket.")
		return
	}

	// Set up the token source.
	const scope = gcs.Scope_ReadOnly
	tokenSrc, err := google.DefaultTokenSource(context.Background(), scope)
	if err != nil {
		err = fmt.Errorf("DefaultTokenSource: %v", err)
		return
	}

	// Use that to create a GCS connection.
	cfg := &gcs.ConnConfig{
		TokenSource: tokenSrc,
	}

	conn, err := gcs.NewConn(cfg)
	if err != nil {
		err = fmt.Errorf("NewConn: %v", err)
		return
	}

	// Open the bucket.
	b, err = conn.OpenBucket(ctx, *fBucket)
	if err != nil {
		err = fmt.Errorf("OpenBucket: %v", err)
		return
	}

	return
}

func getObject(
	ctx context.Context,
	bucket gcs.Bucket) (o *gcs.Object, err error) {
	if *fObject == "" {
		err = errors.New("You must set --object.")
		return
	}

	o, err = bucket.StatObject(
		ctx,
		&gcs.StatObjectRequest{Name: *fObject})

	return
}

// Run workers until SIGINT is received. Return a slice of results.
func runWorkers(
	ctx context.Context,
	o *gcs.Object,
	bucket gcs.Bucket) (results []result, err error) {
	b := syncutil.NewBundle(ctx)

	// Set up a channel that is closed upon SIGINT.
	stop := make(chan struct{})
	go func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, os.Interrupt)

		<-c
		log.Printf("SIGINT received. Stopping soon...")
		close(stop)
	}()

	// Start several workers making random reads.
	var wg sync.WaitGroup
	resultChan := make(chan result)
	for i := 0; i < *fWorkers; i++ {
		wg.Add(1)
		b.Add(func(ctx context.Context) (err error) {
			defer wg.Done()
			err = makeReads(ctx, o, bucket, resultChan, stop)
			if err != nil {
				err = fmt.Errorf("makeReads: %v", err)
				return
			}

			return
		})
	}

	go func() {
		wg.Wait()
		close(resultChan)
	}()

	// Accumulate results.
	b.Add(func(ctx context.Context) (err error) {
		for r := range resultChan {
			results = append(results, r)
		}

		return
	})

	err = b.Join()
	return
}

func describeResults(results []result)

////////////////////////////////////////////////////////////////////////
// Main
////////////////////////////////////////////////////////////////////////

func run(ctx context.Context) (err error) {
	// Open the bucket.
	bucket, err := getBucket(ctx)
	if err != nil {
		err = fmt.Errorf("getBucket: %v", err)
		return
	}

	// Stat the object.
	o, err := getObject(ctx, bucket)
	if err != nil {
		err = fmt.Errorf("getObject: %v", err)
		return
	}

	// Make reads.
	results, err := runWorkers(ctx, o, bucket)
	if err != nil {
		err = fmt.Errorf("runWorkers: %v", err)
		return
	}

	// Print information about the results.
	describeResults(results)

	return
}

func main() {
	log.SetFlags(log.Lmicroseconds)

	err := run(context.Background())
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>describeResults<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// A program that makes random reads within a large object on GCS, reporting
// throughput and latency information.
package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"os/signal"
	"sync"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"

	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/syncutil"
)

var fBucket = flag.String("bucket", "", "The GCS bucket from which to read.")
var fObject = flag.String("object", "", "The object within which to read.")
var fSize = flag.Uint64("size", 1<<21, "The size of each read in bytes.")
var fWorkers = flag.Int("workers", 1, "The number of workers to run.")

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type result struct {
	BytesRead int

	// Time taken from making the HTTP request to receiving the first byte of the
	// response body.
	FirstByteLatency time.Duration

	// Time taken from making the HTTP request to receiving the last byte of the
	// response body.
	FullBodyDuration time.Duration
}

func readOnce(
	ctx context.Context,
	o *gcs.Object,
	bucket gcs.Bucket) (r result, err error) {
	// Is the object large enough?
	if o.Size < *fSize {
		err = fmt.Errorf(
			"Object of size %d not large enough for read size %d",
			o.Size,
			*fSize)

		return
	}

	// Set up an appropriate request.
	req := &gcs.ReadObjectRequest{
		Name:       o.Name,
		Generation: o.Generation,
		Range:      &gcs.ByteRange{},
	}

	req.Range.Start = uint64(rand.Int63n(int64(o.Size - *fSize)))
	req.Range.Limit = req.Range.Start + *fSize

	// Create the reader.
	start := time.Now()
	rc, err := bucket.NewReader(ctx, req)
	if err != nil {
		err = fmt.Errorf("NewReader: %v", err)
		return
	}

	defer func() {
		closeErr := rc.Close()
		if err == nil && closeErr != nil {
			err = fmt.Errorf("Close: %v", closeErr)
		}
	}()

	// Measure the time to first byte.
	_, err = rc.Read([]byte{0})
	if err != nil {
		err = fmt.Errorf("Read: %v", err)
		return
	}

	r.FirstByteLatency = time.Since(start)

	// And the time to read everything.
	n, err := io.Copy(ioutil.Discard, rc)
	if err != nil {
		err = fmt.Errorf("Copy: %v", err)
		return
	}

	r.FullBodyDuration = time.Since(start)
	r.BytesRead = int(n + 1)

	return
}

// Make random reads within the given object, writing results to the supplied
// channel. Stop when the stop channel is closed.
func makeReads(
	ctx context.Context,
	o *gcs.Object,
	bucket gcs.Bucket,
	results chan<- result,
	stop <-chan struct{}) (err error) {
	// Is the object large enough?
	if o.Size < *fSize {
		err = fmt.Errorf(
			"Object of size %d not large enough for read size %d",
			o.Size,
			*fSize)

		return
	}

	for {
		// Stop?
		if _, ok := <-stop; ok {
			return
		}

		// Make a read.
		var r result
		r, err = readOnce(ctx, o, bucket)
		if err != nil {
			err = fmt.Errorf("readOnce %v", err)
			return
		}

		log.Printf(
			"Read %d bytes in %v (first byte latency %v)",
			r.BytesRead,
			r.FullBodyDuration,
			r.FirstByteLatency)

		// Write out the result.
		select {
		case <-ctx.Done():
			err = ctx.Err()
			return

		case results <- r:
		}
	}
}

func getBucket(ctx context.Context) (b gcs.Bucket, err error) {
	if *fBucket == "" {
		err = errors.New("You must set --bucket.")
		return
	}

	// Set up the token source.
	const scope = gcs.Scope_ReadOnly
	tokenSrc, err := google.DefaultTokenSource(context.Background(), scope)
	if err != nil {
		err = fmt.Errorf("DefaultTokenSource: %v", err)
		return
	}

	// Use that to create a GCS connection.
	cfg := &gcs.ConnConfig{
		TokenSource: tokenSrc,
	}

	conn, err := gcs.NewConn(cfg)
	if err != nil {
		err = fmt.Errorf("NewConn: %v", err)
		return
	}

	// Open the bucket.
	b, err = conn.OpenBucket(ctx, *fBucket)
	if err != nil {
		err = fmt.Errorf("OpenBucket: %v", err)
		return
	}

	return
}

func getObject(
	ctx context.Context,
	bucket gcs.Bucket) (o *gcs.Object, err error) {
	if *fObject == "" {
		err = errors.New("You must set --object.")
		return
	}

	o, err = bucket.StatObject(
		ctx,
		&gcs.StatObjectRequest{Name: *fObject})

	return
}

// Run workers until SIGINT is received. Return a slice of results.
func runWorkers(
	ctx context.Context,
	o *gcs.Object,
	bucket gcs.Bucket) (results []result, err error) {
	b := syncutil.NewBundle(ctx)

	// Set up a channel that is closed upon SIGINT.
	stop := make(chan struct{})
	go func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, os.Interrupt)

		<-c
		log.Printf("SIGINT received. Stopping soon...")
		close(stop)
	}()

	// Start several workers making random reads.
	var wg sync.WaitGroup
	resultChan := make(chan result)
	for i := 0; i < *fWorkers; i++ {
		wg.Add(1)
		b.Add(func(ctx context.Context) (err error) {
			defer wg.Done()
			err = makeReads(ctx, o, bucket, resultChan, stop)
			if err != nil {
				err = fmt.Errorf("makeReads: %v", err)
				return
			}

			return
		})
	}

	go func() {
		wg.Wait()
		close(resultChan)
	}()

	// Accumulate results.
	b.Add(func(ctx context.Context) (err error) {
		for r := range resultChan {
			results = append(results, r)
		}

		return
	})

	err = b.Join()
	return
}

func describeResults(results []result) {
	fmt.Printf("Made %d reads.\n", len(results))
}

////////////////////////////////////////////////////////////////////////
// Main
////////////////////////////////////////////////////////////////////////

func run(ctx context.Context) (err error) {
	// Open the bucket.
	bucket, err := getBucket(ctx)
	if err != nil {
		err = fmt.Errorf("getBucket: %v", err)
		return
	}

	// Stat the object.
	o, err := getObject(ctx, bucket)
	if err != nil {
		err = fmt.Errorf("getObject: %v", err)
		return
	}

	// Make reads.
	results, err := runWorkers(ctx, o, bucket)
	if err != nil {
		err = fmt.Errorf("runWorkers: %v", err)
		return
	}

	// Print information about the results.
	describeResults(results)

	return
}

func main() {
	log.SetFlags(log.Lmicroseconds)

	err := run(context.Background())
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package main

// scanning an HTTP response for phrases

import (
	"http"
	"mahonia.googlecode.com/hg"
	"strings"
)

// phrasesInResponse scans the content of an http.Response for phrases,
// and returns a map of phrases and counts.
func phrasesInResponse(res *http.Response) map[string]int {
	defer res.Body.Close()

	contentType := res.Header.Get("Content-Type")

	wr := newWordReader(res.Body, decoderForContentType(contentType))
	ps := newPhraseScanner()
	ps.scanByte(' ')
	buf := make([]byte, 4096)
	for {
		n, err := wr.Read(buf)
		if err != nil {
			break
		}
		for _, c := range buf[:n] {
			ps.scanByte(c)
		}
	}
	ps.scanByte(' ')

	return ps.tally
}

func decoderForContentType(t string) mahonia.Decoder {
	t = strings.ToLower(t)
	var result mahonia.Decoder

	i := strings.Index(t, "charset=")
	if i != -1 {
		charset := t[i+len("charset="):]
		i = strings.Index(charset, ";")
		if i != -1 {
			charset = charset[:i]
		}
		result = mahonia.NewDecoder(charset)
	}

	if result == nil {
		result = mahonia.FallbackDecoder(mahonia.NewDecoder("UTF-8"), mahonia.NewDecoder("windows-1252"))
	}

	if strings.Contains(t, "html") {
		result = mahonia.FallbackDecoder(mahonia.EntityDecoder(), result)
	}

	return result
}
<commit_msg>Log unsupported character sets.<commit_after>package main

// scanning an HTTP response for phrases

import (
	"http"
	"log"
	"mahonia.googlecode.com/hg"
	"strings"
)

// phrasesInResponse scans the content of an http.Response for phrases,
// and returns a map of phrases and counts.
func phrasesInResponse(res *http.Response) map[string]int {
	defer res.Body.Close()

	contentType := res.Header.Get("Content-Type")

	wr := newWordReader(res.Body, decoderForContentType(contentType))
	ps := newPhraseScanner()
	ps.scanByte(' ')
	buf := make([]byte, 4096)
	for {
		n, err := wr.Read(buf)
		if err != nil {
			break
		}
		for _, c := range buf[:n] {
			ps.scanByte(c)
		}
	}
	ps.scanByte(' ')

	return ps.tally
}

func decoderForContentType(t string) mahonia.Decoder {
	t = strings.ToLower(t)
	var result mahonia.Decoder

	i := strings.Index(t, "charset=")
	if i != -1 {
		charset := t[i+len("charset="):]
		i = strings.Index(charset, ";")
		if i != -1 {
			charset = charset[:i]
		}
		result = mahonia.NewDecoder(charset)
		if result == nil {
			log.Println("Unknown charset:", charset)
		}
	}

	if result == nil {
		result = mahonia.FallbackDecoder(mahonia.NewDecoder("UTF-8"), mahonia.NewDecoder("windows-1252"))
	}

	if strings.Contains(t, "html") {
		result = mahonia.FallbackDecoder(mahonia.EntityDecoder(), result)
	}

	return result
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
)

var (
	bindAddress = flag.String("address", ":7701", "address to listen on for metrics (Carbon-compatible protocol)")
	httpAddress = flag.String("http-address", ":7702", "address to recieve queries (http)")
	runAudits = flag.Bool("audit", false, "run audits periodically")
	persist = flag.Bool("persist", false, "persist to disk (load at startup, save on SIGTERM/SIGINT) (see --persist-path)")
	persistPath = flag.String("persist-path", "almaz.dat", "path to storage file")
	persistInterval = flag.Int("bgsave", 0, "save to disk every N seconds (0 --- do not save). Must have --persist specified.")
	debug = flag.Bool("debug", false, "print additional info")
	storageDuration = flag.Int("duration-in-hours", 24, "store metrics for last N hours")
	storagePrecision = flag.Int("precision-in-seconds", 60, "store metrics with precision of N seconds")
	acceptanceRegex = flag.String("regex", "", "accept only metrics which match regular expression")
)

func main() {
	flag.Parse()
	server := NewAlmazServer(*persistPath)
	server.storage.SetStorageParams(*storageDuration, *storagePrecision)
	if *acceptanceRegex != "" {
		server.AddAcceptanceRegex(*acceptanceRegex)
	}
	if *persist {
		server.LoadFromDisk()
	}
	if *runAudits {
		go server.AuditLoop()
	}
	go server.StartGraphite(*bindAddress)
	go server.StartHttpface(*httpAddress)
	server.WaitForTermination(*persist, *persistInterval)
}
<commit_msg>Profiling boilerplate<commit_after>package main

import (
	"flag"
	"log"
	"os"
	"runtime/pprof"
)

var (
	bindAddress      = flag.String("address", ":7701", "address to listen on for metrics (Carbon-compatible protocol)")
	httpAddress      = flag.String("http-address", ":7702", "address to recieve queries (http)")
	runAudits        = flag.Bool("audit", false, "run audits periodically")
	persist          = flag.Bool("persist", false, "persist to disk (load at startup, save on SIGTERM/SIGINT) (see --persist-path)")
	persistPath      = flag.String("persist-path", "almaz.dat", "path to storage file")
	persistInterval  = flag.Int("bgsave", 0, "save to disk every N seconds (0 --- do not save). Must have --persist specified.")
	debug            = flag.Bool("debug", false, "print additional info")
	storageDuration  = flag.Int("duration-in-hours", 24, "store metrics for last N hours")
	storagePrecision = flag.Int("precision-in-seconds", 60, "store metrics with precision of N seconds")
	acceptanceRegex  = flag.String("regex", "", "accept only metrics which match regular expression")
	cpuprofile       = flag.String("cpuprofile", "", "Write cpuprofile info to file")
)

func main() {
	flag.Parse()
	server := NewAlmazServer(*persistPath)
	server.storage.SetStorageParams(*storageDuration, *storagePrecision)
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Println(err.Error())
		}
		log.Printf("Writing cpuprofile to %v\n", *cpuprofile)
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}
	if *acceptanceRegex != "" {
		server.AddAcceptanceRegex(*acceptanceRegex)
	}
	if *persist {
		server.LoadFromDisk()
	}
	if *runAudits {
		go server.AuditLoop()
	}
	go server.StartGraphite(*bindAddress)
	go server.StartHttpface(*httpAddress)
	server.WaitForTermination(*persist, *persistInterval)
}
<|endoftext|>


<commit_before>package store

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"testing"

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

type Howdy struct {
	Yall int `json:"yall"`
}

func ExampleGetMulti() {
	b := NewMemStore()
	b.Set("/howdy/1", []byte(`{"yall": 1}`), 0)
	b.Set("/howdy/2", []byte(`{"yall": 2}`), 0)

	readers := b.GetMulti("/howdy")
	for {
		var h Howdy
		err := json.NewDecoder(readers).Decode(&h)
		if err != nil {
			break
		}
		fmt.Printf("%v\n", h.Yall)
	}
}

func TestMemBackend(t *testing.T) {
	s := NewMemStore()
	assert.Implements(t, (*Store)(nil), new(MemStore), "MemStore")
	testImplementation(t, s)
}

func testImplementation(t *testing.T, s Store) {
	testGet(t, s)
	testGetMulti(t, s)
	testGetMultiMap(t, s)
	testSet(t, s)
	testKeys(t, s)
	testEmptyKeys(t, s)
}

func testGet(t *testing.T, store Store) {
	deleteKeys(store, "/howdy")
	key := "/howdy"
	value := "doody"
	store.Set(key, []byte(value), 0)
	reader := store.Get(key)
	actual := streamToString(reader)
	expected := value
	assert.Equal(t, expected, actual)
}

func testGetMulti(t *testing.T, store Store) {
	deleteKeys(store, "/howdy")
	deleteKeys(store, "/doody")
	data := [3][2]string{{"/howdy/1", `{"yall": 1}`}, {"/howdy/2", `{"yall": 2}`}, {"/doody/1", `{"yall": 1}`}}

	for _, d := range data {
		key := d[0]
		value := d[1]
		store.Set(key, []byte(value), 0)
	}
	readers := store.GetMulti("/howdy")
	for i := 0; i < 2; i++ {
		var h Howdy
		err := json.NewDecoder(readers).Decode(&h)
		assert.NoError(t, err)
		assert.InDelta(t, h.Yall, 1, 1)
	}
	var h Howdy
	err := json.NewDecoder(readers).Decode(&h)
	assert.Error(t, err, "EOF")
}

func testGetMultiMap(t *testing.T, store Store) {
	deleteKeys(store, "/howdy")
	deleteKeys(store, "/doody")
	data := [3][2]string{{"/howdy/1", `{"yall": 1}`}, {"/howdy/2", `{"yall": 2}`}, {"/doody/1", `{"yall": 1}`}}

	for _, d := range data {
		key := d[0]
		value := d[1]
		store.Set(key, []byte(value), 0)
	}
	res, err := store.GetMultiMap("/howdy")
	assert.NoError(t, err)
	assert.Len(t, res, 2)
	val, ok := res["/howdy/1"]
	assert.True(t, ok)
	assert.Equal(t, streamToString(val), `{"yall": 1}`)

	val, ok = res["/howdy/2"]
	assert.True(t, ok)
	assert.Equal(t, streamToString(val), `{"yall": 2}`)
}

func testSet(t *testing.T, store Store) {
	deleteKeys(store, "/howdy")
	value := "value"
	err := store.Set("/howdy", []byte(value), 0)
	assert.NoError(t, err)
	expected := value
	actual := streamToString(store.Get("/howdy"))
	assert.Equal(t, expected, actual)
}

func testEmptyKeys(t *testing.T, store Store) {
	deleteKeys(store, "/howdy")
	actual, err := store.Keys("/howdy")
	assert.Error(t, err)
	assert.Empty(t, actual)
}

func testKeys(t *testing.T, store Store) {
	deleteKeys(store, "/howdy")
	deleteKeys(store, "/foo")
	store.Set("/howdy/doody", []byte("value"), 0)
	store.Set("/howdy/yall", []byte("value"), 0)
	store.Set("/foo", []byte("value"), 0)

	actual, err := store.Keys("/howdy")
	assert.NoError(t, err)
	var expected = []string{"/howdy/doody", "/howdy/yall"}
	assert.Len(t, expected, 2)
	for _, e := range expected {
		assert.Contains(t, actual, e)
	}
}

func deleteKeys(s Store, key string) {
	// es, ok := s.(*EtcdStore)
	// if ok {
	// 	c, q := context.WithTimeout(context.Background(), es.Timeout)
	// 	defer q()
	// 	es.KeysAPI.Delete(c, key, &client.DeleteOptions{Recursive: true})
	// }
	// ms, ok := s.(*MemStore)
	// if ok {
	// 	ms.store = make(map[string]string)
	// }
	s.Delete(key, true)
}

func streamToString(stream io.Reader) string {
	buf := new(bytes.Buffer)
	buf.ReadFrom(stream)
	return buf.String()
}
<commit_msg>clean up<commit_after>package store

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"testing"

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

type Howdy struct {
	Yall int `json:"yall"`
}

func ExampleGetMulti() {
	b := NewMemStore()
	b.Set("/howdy/1", []byte(`{"yall": 1}`), 0)
	b.Set("/howdy/2", []byte(`{"yall": 2}`), 0)

	readers := b.GetMulti("/howdy")
	for {
		var h Howdy
		err := json.NewDecoder(readers).Decode(&h)
		if err != nil {
			break
		}
		fmt.Printf("%v\n", h.Yall)
	}
}

func TestMemBackend(t *testing.T) {
	s := NewMemStore()
	assert.Implements(t, (*Store)(nil), new(MemStore), "MemStore")
	testImplementation(t, s)
}

func testImplementation(t *testing.T, s Store) {
	testGet(t, s)
	testGetMulti(t, s)
	testGetMultiMap(t, s)
	testSet(t, s)
	testKeys(t, s)
	testEmptyKeys(t, s)
}

func testGet(t *testing.T, store Store) {
	deleteKeys(store, "/howdy")
	key := "/howdy"
	value := "doody"
	store.Set(key, []byte(value), 0)
	reader := store.Get(key)
	actual := streamToString(reader)
	expected := value
	assert.Equal(t, expected, actual)
}

func testGetMulti(t *testing.T, store Store) {
	deleteKeys(store, "/howdy")
	deleteKeys(store, "/doody")
	data := [3][2]string{{"/howdy/1", `{"yall": 1}`}, {"/howdy/2", `{"yall": 2}`}, {"/doody/1", `{"yall": 1}`}}

	for _, d := range data {
		key := d[0]
		value := d[1]
		store.Set(key, []byte(value), 0)
	}
	readers := store.GetMulti("/howdy")
	for i := 0; i < 2; i++ {
		var h Howdy
		err := json.NewDecoder(readers).Decode(&h)
		assert.NoError(t, err)
		assert.InDelta(t, h.Yall, 1, 1)
	}
	var h Howdy
	err := json.NewDecoder(readers).Decode(&h)
	assert.Error(t, err, "EOF")
}

func testGetMultiMap(t *testing.T, store Store) {
	deleteKeys(store, "/howdy")
	deleteKeys(store, "/doody")
	data := [3][2]string{{"/howdy/1", `{"yall": 1}`}, {"/howdy/2", `{"yall": 2}`}, {"/doody/1", `{"yall": 1}`}}

	for _, d := range data {
		key := d[0]
		value := d[1]
		store.Set(key, []byte(value), 0)
	}
	res, err := store.GetMultiMap("/howdy")
	assert.NoError(t, err)
	assert.Len(t, res, 2)
	val, ok := res["/howdy/1"]
	assert.True(t, ok)
	assert.Equal(t, streamToString(val), `{"yall": 1}`)

	val, ok = res["/howdy/2"]
	assert.True(t, ok)
	assert.Equal(t, streamToString(val), `{"yall": 2}`)
}

func testSet(t *testing.T, store Store) {
	deleteKeys(store, "/howdy")
	value := "value"
	err := store.Set("/howdy", []byte(value), 0)
	assert.NoError(t, err)
	expected := value
	actual := streamToString(store.Get("/howdy"))
	assert.Equal(t, expected, actual)
}

func testEmptyKeys(t *testing.T, store Store) {
	deleteKeys(store, "/howdy")
	actual, err := store.Keys("/howdy")
	assert.Error(t, err)
	assert.Empty(t, actual)
}

func testKeys(t *testing.T, store Store) {
	deleteKeys(store, "/howdy")
	deleteKeys(store, "/foo")
	store.Set("/howdy/doody", []byte("value"), 0)
	store.Set("/howdy/yall", []byte("value"), 0)
	store.Set("/foo", []byte("value"), 0)

	actual, err := store.Keys("/howdy")
	assert.NoError(t, err)
	var expected = []string{"/howdy/doody", "/howdy/yall"}
	assert.Len(t, expected, 2)
	for _, e := range expected {
		assert.Contains(t, actual, e)
	}
}

func deleteKeys(s Store, key string) {
	s.Delete(key, true)
}

func streamToString(stream io.Reader) string {
	buf := new(bytes.Buffer)
	buf.ReadFrom(stream)
	return buf.String()
}
<|endoftext|>


<commit_before>package strava

import (
	"log"
	"time"

	stravalib "github.com/svdberg/syncmysport-runkeeper/Godeps/_workspace/src/github.com/strava/go.strava"
	dm "github.com/svdberg/syncmysport-runkeeper/datamodel"
)

func ConvertToActivity(stravaActivity *stravalib.ActivityDetailed, timeStream *stravalib.StreamSet, gpsTrack *stravalib.StreamSet, hrTrack *stravalib.StreamSet, altTrack *stravalib.StreamSet) *dm.Activity {
	stvActivity := dm.CreateActivity()
	stvActivity.StartTime = int(stravaActivity.StartDate.Unix()) //UTC date
	log.Printf("STV Local date: %s, start date: %s, unix: %d", stravaActivity.StartDateLocal, stravaActivity.StartDate, stravaActivity.StartDate.Unix())
	stvActivity.Duration = stravaActivity.ElapsedTime
	stvActivity.Name = stravaActivity.Name
	stvActivity.Calories = stravaActivity.Calories
	stvActivity.Distance = stravaActivity.Distance
	stvActivity.AverageHeartRate = int(stravaActivity.AverageHeartrate)
	loc, err := time.LoadLocation(stravaActivity.TimeZone)
	if err != nil {
		timeInTZ := time.Time(stravaActivity.StartDate).In(loc)
		_, offsetInSeconds := timeInTZ.Zone()
		stvActivity.UtcOffSet = offsetInSeconds / 60 / 60
	} else {
		log.Printf("Warning: reading location from strava Activity failed with: %s", err)
	}

	if stravaActivity.Type.String() == "Run" {
		stvActivity.Type = "Running"
	} else if stravaActivity.Type.String() == "Ride" || stravaActivity.Type.String() == "EBikeRide" || stravaActivity.Type.String() == "VirtualRide" {
		stvActivity.Type = "Cycling"
	} else if stravaActivity.Type.String() == "Swim" {
		stvActivity.Type = "Swimming"
	} else {
		//I don't know, call it Activity
		stvActivity.Type = "Activity"
	}

	if gpsTrack != nil && gpsTrack.Location != nil && timeStream != nil {
		stvActivity.GPS = convertGPSTrack(gpsTrack, timeStream, altTrack)
	}
	if hrTrack != nil && hrTrack.HeartRate != nil && timeStream != nil {
		stvActivity.HeartRate = convertHeartRateTrack(hrTrack, timeStream)
	}
	return stvActivity
}

func convertGPSTrack(sourceStream *stravalib.StreamSet, timeStream *stravalib.StreamSet, elevationStream *stravalib.StreamSet) []dm.GPS {
	//merge the time stream + the location stream
	merged := mergeTimeAndLocation(timeStream.Time, sourceStream.Location, elevationStream.Elevation)

	result := make([]dm.GPS, len(sourceStream.Location.Data))
	for index, gpsTime := range merged {
		result[index] = dm.GPS{float64(gpsTime.Time), gpsTime.Alt, gpsTime.Long, gpsTime.Lat}
	}
	return result
}

type GPSTime struct {
	Time int
	Lat  float64
	Long float64
	Alt  float64
}

func mergeTimeAndLocation(timeStream *stravalib.IntegerStream, locStream *stravalib.LocationStream, altStream *stravalib.DecimalStream) []GPSTime {
	merged := make([]GPSTime, len(timeStream.Data))
	for i, t := range timeStream.Data {
		latLong := locStream.Data[i]
		alt := altStream.Data[i]
		merged[i] = GPSTime{t, latLong[0], latLong[1], alt}
	}
	return merged
}

func convertHeartRateTrack(sourceStream *stravalib.StreamSet, timeStream *stravalib.StreamSet) []dm.HeartRate {
	result := make([]dm.HeartRate, len(sourceStream.HeartRate.Data))
	for index, hr := range sourceStream.HeartRate.Data {
		hrTime := timeStream.Time.Data[index]
		result[index] = dm.HeartRate{float64(hrTime), hr}
	}
	return result
}
<commit_msg>Fix error printing<commit_after>package strava

import (
	"log"
	"time"

	stravalib "github.com/svdberg/syncmysport-runkeeper/Godeps/_workspace/src/github.com/strava/go.strava"
	dm "github.com/svdberg/syncmysport-runkeeper/datamodel"
)

func ConvertToActivity(stravaActivity *stravalib.ActivityDetailed, timeStream *stravalib.StreamSet, gpsTrack *stravalib.StreamSet, hrTrack *stravalib.StreamSet, altTrack *stravalib.StreamSet) *dm.Activity {
	stvActivity := dm.CreateActivity()
	stvActivity.StartTime = int(stravaActivity.StartDate.Unix()) //UTC date
	log.Printf("STV Local date: %s, start date: %s, unix: %d", stravaActivity.StartDateLocal, stravaActivity.StartDate, stravaActivity.StartDate.Unix())
	stvActivity.Duration = stravaActivity.ElapsedTime
	stvActivity.Name = stravaActivity.Name
	stvActivity.Calories = stravaActivity.Calories
	stvActivity.Distance = stravaActivity.Distance
	stvActivity.AverageHeartRate = int(stravaActivity.AverageHeartrate)
	loc, err := time.LoadLocation(stravaActivity.TimeZone)
	if err != nil {
		timeInTZ := time.Time(stravaActivity.StartDate).In(loc)
		_, offsetInSeconds := timeInTZ.Zone()
		stvActivity.UtcOffSet = offsetInSeconds / 60 / 60
	} else {
		log.Printf("Warning: reading location from strava Activity failed with: %e", err)
	}

	if stravaActivity.Type.String() == "Run" {
		stvActivity.Type = "Running"
	} else if stravaActivity.Type.String() == "Ride" || stravaActivity.Type.String() == "EBikeRide" || stravaActivity.Type.String() == "VirtualRide" {
		stvActivity.Type = "Cycling"
	} else if stravaActivity.Type.String() == "Swim" {
		stvActivity.Type = "Swimming"
	} else {
		//I don't know, call it Activity
		stvActivity.Type = "Activity"
	}

	if gpsTrack != nil && gpsTrack.Location != nil && timeStream != nil {
		stvActivity.GPS = convertGPSTrack(gpsTrack, timeStream, altTrack)
	}
	if hrTrack != nil && hrTrack.HeartRate != nil && timeStream != nil {
		stvActivity.HeartRate = convertHeartRateTrack(hrTrack, timeStream)
	}
	return stvActivity
}

func convertGPSTrack(sourceStream *stravalib.StreamSet, timeStream *stravalib.StreamSet, elevationStream *stravalib.StreamSet) []dm.GPS {
	//merge the time stream + the location stream
	merged := mergeTimeAndLocation(timeStream.Time, sourceStream.Location, elevationStream.Elevation)

	result := make([]dm.GPS, len(sourceStream.Location.Data))
	for index, gpsTime := range merged {
		result[index] = dm.GPS{float64(gpsTime.Time), gpsTime.Alt, gpsTime.Long, gpsTime.Lat}
	}
	return result
}

type GPSTime struct {
	Time int
	Lat  float64
	Long float64
	Alt  float64
}

func mergeTimeAndLocation(timeStream *stravalib.IntegerStream, locStream *stravalib.LocationStream, altStream *stravalib.DecimalStream) []GPSTime {
	merged := make([]GPSTime, len(timeStream.Data))
	for i, t := range timeStream.Data {
		latLong := locStream.Data[i]
		alt := altStream.Data[i]
		merged[i] = GPSTime{t, latLong[0], latLong[1], alt}
	}
	return merged
}

func convertHeartRateTrack(sourceStream *stravalib.StreamSet, timeStream *stravalib.StreamSet) []dm.HeartRate {
	result := make([]dm.HeartRate, len(sourceStream.HeartRate.Data))
	for index, hr := range sourceStream.HeartRate.Data {
		hrTime := timeStream.Time.Data[index]
		result[index] = dm.HeartRate{float64(hrTime), hr}
	}
	return result
}
<|endoftext|>


<commit_before>// This file is adapted from code in the net/http/httputil
// package of the Go standard library, which is by the
// Go Authors, and bears this copyright and license info:
//
//   Copyright 2011 The Go Authors. All rights reserved.
//   Use of this source code is governed by a BSD-style
//   license that can be found in the LICENSE file.
//
// This file has been modified from the standard lib to
// meet the needs of the application.

package proxy

import (
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

// onExitFlushLoop is a callback set by tests to detect the state of the
// flushLoop() goroutine.
var onExitFlushLoop func()

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// Director must be a function which modifies
	// the request into a new request to be sent
	// using Transport. Its response is then copied
	// back to the original client unmodified.
	Director func(*http.Request)

	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration
}

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

// NewSingleHostReverseProxy returns a new ReverseProxy that rewrites
// URLs to the scheme, host, and base path provided in target. If the
// target's path is "/base" and the incoming request was for "/dir",
// the target request will be for /base/dir.
// Without logic: target's path is "/", incoming is "/api/messages",
// without is "/api", then the target request will be for /messages.
func NewSingleHostReverseProxy(target *url.URL, without string) *ReverseProxy {
	targetQuery := target.RawQuery
	director := func(req *http.Request) {
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		req.URL.Path = singleJoiningSlash(target.Path, req.URL.Path)
		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}
		if without != "" {
			req.URL.Path = strings.TrimPrefix(req.URL.Path, without)
		}
	}
	return &ReverseProxy{Director: director}
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request, extraHeaders http.Header) error {
	transport := p.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}

	outreq := new(http.Request)
	*outreq = *req // includes shallow copies of maps, but okay

	p.Director(outreq)
	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	// Remove hop-by-hop headers to the backend.  Especially
	// important is "Connection" because we want a persistent
	// connection, regardless of what the client sent to us.  This
	// is modifying the same underlying map from req (shallow
	// copied above) so we only copy it if necessary.
	copiedHeaders := false
	for _, h := range hopHeaders {
		if outreq.Header.Get(h) != "" {
			if !copiedHeaders {
				outreq.Header = make(http.Header)
				copyHeader(outreq.Header, req.Header)
				copiedHeaders = true
			}
			outreq.Header.Del(h)
		}
	}

	if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
		// If we aren't the first proxy retain prior
		// X-Forwarded-For information as a comma+space
		// separated list and fold multiple headers into one.
		if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		outreq.Header.Set("X-Forwarded-For", clientIP)
	}

	if extraHeaders != nil {
		for k, v := range extraHeaders {
			outreq.Header[k] = v
		}
	}

	res, err := transport.RoundTrip(outreq)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode == http.StatusSwitchingProtocols && res.Header.Get("Upgrade") == "websocket" {
		hj, ok := rw.(http.Hijacker)
		if !ok {
			return nil
		}

		conn, _, err := hj.Hijack()
		if err != nil {
			return err
		}
		defer conn.Close()

		backendConn, err := net.Dial("tcp", outreq.Host)
		if err != nil {
			return err
		}
		defer backendConn.Close()

		outreq.Write(backendConn)

		go func() {
			io.Copy(backendConn, conn) // write tcp stream to backend.
		}()
		io.Copy(conn, backendConn) // read tcp stream from backend.
	} else {
		for _, h := range hopHeaders {
			res.Header.Del(h)
		}

		copyHeader(rw.Header(), res.Header)

		rw.WriteHeader(res.StatusCode)
		p.copyResponse(rw, res.Body)
	}

	return nil
}

func (p *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	if p.FlushInterval != 0 {
		if wf, ok := dst.(writeFlusher); ok {
			mlw := &maxLatencyWriter{
				dst:     wf,
				latency: p.FlushInterval,
				done:    make(chan bool),
			}
			go mlw.flushLoop()
			defer mlw.stop()
			dst = mlw
		}
	}

	io.Copy(dst, src)
}

type writeFlusher interface {
	io.Writer
	http.Flusher
}

type maxLatencyWriter struct {
	dst     writeFlusher
	latency time.Duration

	lk   sync.Mutex // protects Write + Flush
	done chan bool
}

func (m *maxLatencyWriter) Write(p []byte) (int, error) {
	m.lk.Lock()
	defer m.lk.Unlock()
	return m.dst.Write(p)
}

func (m *maxLatencyWriter) flushLoop() {
	t := time.NewTicker(m.latency)
	defer t.Stop()
	for {
		select {
		case <-m.done:
			if onExitFlushLoop != nil {
				onExitFlushLoop()
			}
			return
		case <-t.C:
			m.lk.Lock()
			m.dst.Flush()
			m.lk.Unlock()
		}
	}
}

func (m *maxLatencyWriter) stop() { m.done <- true }
<commit_msg>Set default MaxIdleConnsPerHost to 64<commit_after>// This file is adapted from code in the net/http/httputil
// package of the Go standard library, which is by the
// Go Authors, and bears this copyright and license info:
//
//   Copyright 2011 The Go Authors. All rights reserved.
//   Use of this source code is governed by a BSD-style
//   license that can be found in the LICENSE file.
//
// This file has been modified from the standard lib to
// meet the needs of the application.

package proxy

import (
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

// onExitFlushLoop is a callback set by tests to detect the state of the
// flushLoop() goroutine.
var onExitFlushLoop func()

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// Director must be a function which modifies
	// the request into a new request to be sent
	// using Transport. Its response is then copied
	// back to the original client unmodified.
	Director func(*http.Request)

	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration
}

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

// NewSingleHostReverseProxy returns a new ReverseProxy that rewrites
// URLs to the scheme, host, and base path provided in target. If the
// target's path is "/base" and the incoming request was for "/dir",
// the target request will be for /base/dir.
// Without logic: target's path is "/", incoming is "/api/messages",
// without is "/api", then the target request will be for /messages.
func NewSingleHostReverseProxy(target *url.URL, without string) *ReverseProxy {
	targetQuery := target.RawQuery
	director := func(req *http.Request) {
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		req.URL.Path = singleJoiningSlash(target.Path, req.URL.Path)
		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}
		if without != "" {
			req.URL.Path = strings.TrimPrefix(req.URL.Path, without)
		}
	}
	return &ReverseProxy{Director: director}
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

var DefaultTransport http.RoundTripper = &http.Transport{
	Proxy: http.ProxyFromEnvironment,
	Dial: (&net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
	}).Dial,
	TLSHandshakeTimeout: 10 * time.Second,
	MaxIdleConnsPerHost: 64,
}

func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request, extraHeaders http.Header) error {
	transport := p.Transport
	if transport == nil {
		transport = DefaultTransport
	}

	outreq := new(http.Request)
	*outreq = *req // includes shallow copies of maps, but okay

	p.Director(outreq)
	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	// Remove hop-by-hop headers to the backend.  Especially
	// important is "Connection" because we want a persistent
	// connection, regardless of what the client sent to us.  This
	// is modifying the same underlying map from req (shallow
	// copied above) so we only copy it if necessary.
	copiedHeaders := false
	for _, h := range hopHeaders {
		if outreq.Header.Get(h) != "" {
			if !copiedHeaders {
				outreq.Header = make(http.Header)
				copyHeader(outreq.Header, req.Header)
				copiedHeaders = true
			}
			outreq.Header.Del(h)
		}
	}

	if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
		// If we aren't the first proxy retain prior
		// X-Forwarded-For information as a comma+space
		// separated list and fold multiple headers into one.
		if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		outreq.Header.Set("X-Forwarded-For", clientIP)
	}

	if extraHeaders != nil {
		for k, v := range extraHeaders {
			outreq.Header[k] = v
		}
	}

	res, err := transport.RoundTrip(outreq)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode == http.StatusSwitchingProtocols && res.Header.Get("Upgrade") == "websocket" {
		hj, ok := rw.(http.Hijacker)
		if !ok {
			return nil
		}

		conn, _, err := hj.Hijack()
		if err != nil {
			return err
		}
		defer conn.Close()

		backendConn, err := net.Dial("tcp", outreq.Host)
		if err != nil {
			return err
		}
		defer backendConn.Close()

		outreq.Write(backendConn)

		go func() {
			io.Copy(backendConn, conn) // write tcp stream to backend.
		}()
		io.Copy(conn, backendConn) // read tcp stream from backend.
	} else {
		for _, h := range hopHeaders {
			res.Header.Del(h)
		}

		copyHeader(rw.Header(), res.Header)

		rw.WriteHeader(res.StatusCode)
		p.copyResponse(rw, res.Body)
	}

	return nil
}

func (p *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	if p.FlushInterval != 0 {
		if wf, ok := dst.(writeFlusher); ok {
			mlw := &maxLatencyWriter{
				dst:     wf,
				latency: p.FlushInterval,
				done:    make(chan bool),
			}
			go mlw.flushLoop()
			defer mlw.stop()
			dst = mlw
		}
	}

	io.Copy(dst, src)
}

type writeFlusher interface {
	io.Writer
	http.Flusher
}

type maxLatencyWriter struct {
	dst     writeFlusher
	latency time.Duration

	lk   sync.Mutex // protects Write + Flush
	done chan bool
}

func (m *maxLatencyWriter) Write(p []byte) (int, error) {
	m.lk.Lock()
	defer m.lk.Unlock()
	return m.dst.Write(p)
}

func (m *maxLatencyWriter) flushLoop() {
	t := time.NewTicker(m.latency)
	defer t.Stop()
	for {
		select {
		case <-m.done:
			if onExitFlushLoop != nil {
				onExitFlushLoop()
			}
			return
		case <-t.C:
			m.lk.Lock()
			m.dst.Flush()
			m.lk.Unlock()
		}
	}
}

func (m *maxLatencyWriter) stop() { m.done <- true }
<|endoftext|>


<commit_before><commit_msg>Remove file accidentally added<commit_after><|endoftext|>