<commit_before>package air

import (
	"bytes"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/fsnotify/fsnotify"
	"github.com/tdewolff/minify"
	"github.com/tdewolff/minify/html"
)

type (
	// Renderer is used to provide a `Render()` method for an `Air` instance for renders a
	// "text/html" HTTP response.
	Renderer interface {
		// SetTemplateFunc sets the func f into template func map with the name.
		SetTemplateFunc(name string, f interface{})

		// ParseTemplates parses template files. It will be called in the `Air#Serve()`.
		ParseTemplates() error

		// Render renders the data into the w with the templateName.
		Render(w io.Writer, templateName string, data JSONMap) error
	}

	// renderer implements the `Renderer` by using the `template.Template`.
	renderer struct {
		air *Air

		template        *template.Template
		templateFuncMap template.FuncMap
		minifier        *minify.M
		watcher         *fsnotify.Watcher
	}
)

// newRenderer returns a pointer of a new instance of the `renderer`.
func newRenderer(a *Air) *renderer {
	return &renderer{
		air:      a,
		template: template.New("template"),
		templateFuncMap: template.FuncMap{
			"strlen":  strlen,
			"strcat":  strcat,
			"substr":  substr,
			"timefmt": timefmt,
		},
	}
}

// SetTemplateFunc implements the `Renderer#SetTemplateFunc()` by using the `template.Template`.
func (r *renderer) SetTemplateFunc(name string, f interface{}) {
	r.templateFuncMap[name] = f
}

// ParseTemplates implements the `Renderer#ParseTemplates()` by using the `template.Template`.
//
// e.g. r.air.Config.TemplateRoot == "templates" && r.air.Config.TemplateExt == ".html"
//
// templates/
//   index.html
//   login.html
//   register.html
//
// templates/parts/
//   header.html
//   footer.html
//
// will be parsed into:
//
// "index.html", "login.html", "register.html", "parts/header.html", "parts/footer.html".
func (r *renderer) ParseTemplates() error {
	c := r.air.Config

	if _, err := os.Stat(c.TemplateRoot); err != nil && os.IsNotExist(err) {
		return nil
	}

	if c.TemplateMinified {
		r.minifier = minify.New()
		r.minifier.Add("text/html", &html.Minifier{
			KeepDefaultAttrVals: true,
			KeepDocumentTags:    true,
			KeepWhitespace:      true,
		})
	}

	if c.TemplateWatched {
		var err error
		if r.watcher, err = fsnotify.NewWatcher(); err != nil {
			return err
		}

		dirs, err := walkDirs(r.air.Config.TemplateRoot)
		if err != nil {
			return err
		}

		for _, dir := range dirs {
			if err := r.watcher.Add(dir); err != nil {
				return err
			}
		}

		go r.watchTemplates()
	}

	return r.parseTemplates()
}

// Render implements the `Renderer#Render()` by using the `template.Template`.
func (r *renderer) Render(w io.Writer, templateName string, data JSONMap) error {
	return r.template.ExecuteTemplate(w, templateName, data)
}

// parseTemplates parses all template files.
func (r *renderer) parseTemplates() error {
	c := r.air.Config

	tr := filepath.Clean(c.TemplateRoot)
	if _, err := os.Stat(tr); err != nil && os.IsNotExist(err) {
		return nil
	}

	dirs, err := walkDirs(tr)
	if err != nil {
		return err
	}

	var filenames []string
	for _, dir := range dirs {
		fns, err := filepath.Glob(fmt.Sprintf("%s/*%s", dir, c.TemplateExt))
		if err != nil {
			return err
		}
		filenames = append(filenames, fns...)
	}

	buf := &bytes.Buffer{}

	t := template.New("template")
	t.Funcs(r.templateFuncMap)
	t.Delims(c.TemplateLeftDelim, c.TemplateRightDelim)

	for _, filename := range filenames {
		b, err := ioutil.ReadFile(filename)
		if err != nil {
			return err
		}

		if c.TemplateMinified {
			err := r.minifier.Minify("text/html", buf, bytes.NewReader(b))
			if err != nil {
				return err
			}
			b = buf.Bytes()
			buf.Reset()
		}

		start := 0
		if tr != "." {
			start = len(tr) + 1
		}

		name := filepath.ToSlash(filename[start:])
		if _, err := t.New(name).Parse(string(b)); err != nil {
			return err
		}
	}

	r.template = t

	return nil
}

// watchTemplates watchs the changing of all template files.
func (r *renderer) watchTemplates() {
	for {
		select {
		case event := <-r.watcher.Events:
			r.air.Logger.Info(event)

			if event.Op == fsnotify.Create {
				s := event.String()
				s = s[:strings.Index(s, ":")]
				s = s[1 : len(s)-1]
				if filepath.Ext(s) != r.air.Config.TemplateExt {
					r.watcher.Add(s)
				}
			}

			if err := r.parseTemplates(); err != nil {
				r.air.Logger.Error(err)
			}
		case err := <-r.watcher.Errors:
			r.air.Logger.Error(err)
		}
	}
}

// walkDirs walks all subdirs of the root recursively.
func walkDirs(root string) ([]string, error) {
	var dirs []string
	err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if info.IsDir() {
			dirs = append(dirs, path)
		}
		return err
	})
	return dirs, err
}

// strlen returns the number of chars in the s.
func strlen(s string) int {
	return len([]rune(s))
}

// strcat returns a string that is catenated to the tail of the s by the ss.
func strcat(s string, ss ...string) string {
	for i := range ss {
		s = fmt.Sprintf("%s%s", s, ss[i])
	}
	return s
}

// substr returns the substring consisting of the chars of the s starting at the index i and
// continuing up to, but not including, the char at the index j.
func substr(s string, i, j int) string {
	rs := []rune(s)
	return string(rs[i:j])
}

// timefmt returns a textual representation of the t formatted according to the layout.
func timefmt(t time.Time, layout string) string {
	return t.Format(layout)
}
<commit_msg>refactor: simplify the `renderer#watchTemplates()`<commit_after>package air

import (
	"bytes"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"time"

	"github.com/fsnotify/fsnotify"
	"github.com/tdewolff/minify"
	"github.com/tdewolff/minify/html"
)

type (
	// Renderer is used to provide a `Render()` method for an `Air` instance for renders a
	// "text/html" HTTP response.
	Renderer interface {
		// SetTemplateFunc sets the func f into template func map with the name.
		SetTemplateFunc(name string, f interface{})

		// ParseTemplates parses template files. It will be called in the `Air#Serve()`.
		ParseTemplates() error

		// Render renders the data into the w with the templateName.
		Render(w io.Writer, templateName string, data JSONMap) error
	}

	// renderer implements the `Renderer` by using the `template.Template`.
	renderer struct {
		air *Air

		template        *template.Template
		templateFuncMap template.FuncMap
		minifier        *minify.M
		watcher         *fsnotify.Watcher
	}
)

// newRenderer returns a pointer of a new instance of the `renderer`.
func newRenderer(a *Air) *renderer {
	return &renderer{
		air:      a,
		template: template.New("template"),
		templateFuncMap: template.FuncMap{
			"strlen":  strlen,
			"strcat":  strcat,
			"substr":  substr,
			"timefmt": timefmt,
		},
	}
}

// SetTemplateFunc implements the `Renderer#SetTemplateFunc()` by using the `template.Template`.
func (r *renderer) SetTemplateFunc(name string, f interface{}) {
	r.templateFuncMap[name] = f
}

// ParseTemplates implements the `Renderer#ParseTemplates()` by using the `template.Template`.
//
// e.g. r.air.Config.TemplateRoot == "templates" && r.air.Config.TemplateExt == ".html"
//
// templates/
//   index.html
//   login.html
//   register.html
//
// templates/parts/
//   header.html
//   footer.html
//
// will be parsed into:
//
// "index.html", "login.html", "register.html", "parts/header.html", "parts/footer.html".
func (r *renderer) ParseTemplates() error {
	c := r.air.Config

	if _, err := os.Stat(c.TemplateRoot); err != nil && os.IsNotExist(err) {
		return nil
	}

	if c.TemplateMinified {
		r.minifier = minify.New()
		r.minifier.Add("text/html", &html.Minifier{
			KeepDefaultAttrVals: true,
			KeepDocumentTags:    true,
			KeepWhitespace:      true,
		})
	}

	if c.TemplateWatched {
		var err error
		if r.watcher, err = fsnotify.NewWatcher(); err != nil {
			return err
		}

		dirs, err := walkDirs(r.air.Config.TemplateRoot)
		if err != nil {
			return err
		}

		for _, dir := range dirs {
			if err := r.watcher.Add(dir); err != nil {
				return err
			}
		}

		go r.watchTemplates()
	}

	return r.parseTemplates()
}

// Render implements the `Renderer#Render()` by using the `template.Template`.
func (r *renderer) Render(w io.Writer, templateName string, data JSONMap) error {
	return r.template.ExecuteTemplate(w, templateName, data)
}

// parseTemplates parses all template files.
func (r *renderer) parseTemplates() error {
	c := r.air.Config

	tr := filepath.Clean(c.TemplateRoot)
	if _, err := os.Stat(tr); err != nil && os.IsNotExist(err) {
		return nil
	}

	dirs, err := walkDirs(tr)
	if err != nil {
		return err
	}

	var filenames []string
	for _, dir := range dirs {
		fns, err := filepath.Glob(fmt.Sprintf("%s/*%s", dir, c.TemplateExt))
		if err != nil {
			return err
		}
		filenames = append(filenames, fns...)
	}

	buf := &bytes.Buffer{}

	t := template.New("template")
	t.Funcs(r.templateFuncMap)
	t.Delims(c.TemplateLeftDelim, c.TemplateRightDelim)

	start := 0
	if tr != "." {
		start = len(tr) + 1
	}

	for _, filename := range filenames {
		b, err := ioutil.ReadFile(filename)
		if err != nil {
			return err
		}

		if c.TemplateMinified {
			err := r.minifier.Minify("text/html", buf, bytes.NewReader(b))
			if err != nil {
				return err
			}
			b = buf.Bytes()
			buf.Reset()
		}

		name := filepath.ToSlash(filename[start:])
		if _, err := t.New(name).Parse(string(b)); err != nil {
			return err
		}
	}

	r.template = t

	return nil
}

// watchTemplates watchs the changing of all template files.
func (r *renderer) watchTemplates() {
	for {
		select {
		case event := <-r.watcher.Events:
			r.air.Logger.Info(event)

			if event.Op == fsnotify.Create {
				r.watcher.Add(event.Name)
			}

			if err := r.parseTemplates(); err != nil {
				r.air.Logger.Error(err)
			}
		case err := <-r.watcher.Errors:
			r.air.Logger.Error(err)
		}
	}
}

// walkDirs walks all subdirs of the root recursively.
func walkDirs(root string) ([]string, error) {
	var dirs []string
	err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if info.IsDir() {
			dirs = append(dirs, path)
		}
		return err
	})
	return dirs, err
}

// strlen returns the number of chars in the s.
func strlen(s string) int {
	return len([]rune(s))
}

// strcat returns a string that is catenated to the tail of the s by the ss.
func strcat(s string, ss ...string) string {
	for i := range ss {
		s = fmt.Sprintf("%s%s", s, ss[i])
	}
	return s
}

// substr returns the substring consisting of the chars of the s starting at the index i and
// continuing up to, but not including, the char at the index j.
func substr(s string, i, j int) string {
	rs := []rune(s)
	return string(rs[i:j])
}

// timefmt returns a textual representation of the t formatted according to the layout.
func timefmt(t time.Time, layout string) string {
	return t.Format(layout)
}
<|endoftext|>


<commit_before>// Package analytics provides the methods to run an analytics reporting system
// for API requests which may be useful to users for measuring access and
// possibly identifying bad actors abusing requests.
package analytics

import (
	"encoding/json"
	"log"
	"net/http"
	"runtime"
	"strings"
	"time"

	"github.com/boltdb/bolt"
)

type apiRequest struct {
	URL        string `json:"url"`
	Method     string `json:"http_method"`
	Origin     string `json:"origin"`
	Proto      string `json:"http_protocol"`
	RemoteAddr string `json:"ip_address"`
	Timestamp  int64  `json:"timestamp"`
	External   bool   `json:"external"`
}

var (
	store       *bolt.DB
	requestChan chan apiRequest
)

// Record queues an apiRequest for metrics
func Record(req *http.Request) {
	external := strings.Contains(req.URL.Path, "/external/")

	r := apiRequest{
		URL:        req.URL.String(),
		Method:     req.Method,
		Origin:     req.Header.Get("Origin"),
		Proto:      req.Proto,
		RemoteAddr: req.RemoteAddr,
		Timestamp:  time.Now().Unix() * 1000,
		External:   external,
	}

	// put r on buffered requestChan to take advantage of batch insertion in DB
	requestChan <- r
}

// Close exports the abillity to close our db file. Should be called with defer
// after call to Init() from the same place.
func Close() {
	err := store.Close()
	if err != nil {
		log.Println(err)
	}
}

// Init creates a db connection, initializes the db with schema and data and
// sets up the queue/batching channel
func Init() {
	var err error
	store, err = bolt.Open("analytics.db", 0666, nil)
	if err != nil {
		log.Fatalln(err)
	}

	err = store.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte("__requests"))
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		log.Fatalln("Error idempotently creating requests bucket in analytics.db:", err)
	}

	requestChan = make(chan apiRequest, 1024*64*runtime.NumCPU())

	go serve()

	if err != nil {
		log.Fatalln(err)
	}
}

func serve() {
	// make timer to notify select to batch request insert from requestChan
	// interval: 30 seconds
	apiRequestTimer := time.NewTicker(time.Second * 30)

	// make timer to notify select to remove analytics older than 14 days
	// interval: 1 weeks
	// TODO: enable analytics backup service to cloud
	pruneThreshold := time.Hour * 24 * 14
	pruneDBTimer := time.NewTicker(pruneThreshold / 2)

	for {
		select {
		case <-apiRequestTimer.C:
			err := batchInsert(requestChan)
			if err != nil {
				log.Println(err)
			}

		case <-pruneDBTimer.C:
			err := batchPrune(pruneThreshold)
			if err != nil {
				log.Println(err)
			}

		case <-time.After(time.Second * 30):
			continue
		}
	}
}

// ChartData returns the map containing decoded javascript needed to chart 2 weeks of data by day
func ChartData() (map[string]interface{}, error) {
	// set thresholds for today and the 6 days preceeding
	times := [14]time.Time{}
	dates := [14]string{}
	now := time.Now()
	today := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC)

	for i := range times {
		// subtract 24 * i hours to make days prior
		dur := time.Duration(24 * i * -1)
		day := today.Add(time.Hour * dur)

		// day threshold is [...n-1-i, n-1, n]
		times[len(times)-1-i] = day
		dates[len(times)-1-i] = day.Format("01/02")
	}

	// get api request analytics from db
	var requests = []apiRequest{}
	err := store.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("__requests"))

		err := b.ForEach(func(k, v []byte) error {
			var r apiRequest
			err := json.Unmarshal(v, &r)
			if err != nil {
				log.Println("Error decoding json from analytics db:", err)
				return nil
			}

			requests = append(requests, r)

			return nil
		})
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return nil, err
	}

	ips := [14]map[string]struct{}{}
	for i := range ips {
		ips[i] = make(map[string]struct{})
	}

	total := [14]int{}
	unique := [14]int{}

CHECK_REQUEST:
	for i := range requests {
		ts := time.Unix(requests[i].Timestamp/1000, 0)

		for j := range times {
			// if on today, there will be no next iteration to set values for
			// day prior so all valid requests belong to today
			if j == len(times)-1 {
				if ts.After(times[j]) || ts.Equal(times[j]) {
					// do all record keeping
					total[j]++

					if _, ok := ips[j][requests[i].RemoteAddr]; !ok {
						unique[j]++
						ips[j][requests[i].RemoteAddr] = struct{}{}
					}

					continue CHECK_REQUEST
				}
			}

			if ts.Equal(times[j]) {
				// increment total count for current time threshold (day)
				total[j]++

				// if no IP found for current threshold, increment unique and record IP
				if _, ok := ips[j][requests[i].RemoteAddr]; !ok {
					unique[j]++
					ips[j][requests[i].RemoteAddr] = struct{}{}
				}

				continue CHECK_REQUEST
			}

			if ts.Before(times[j]) {
				// check if older than earliest threshold
				if j == 0 {
					continue CHECK_REQUEST
				}

				// increment total count for previous time threshold (day)
				total[j-1]++

				// if no IP found for day prior, increment unique and record IP
				if _, ok := ips[j-1][requests[i].RemoteAddr]; !ok {
					unique[j-1]++
					ips[j-1][requests[i].RemoteAddr] = struct{}{}
				}
			}
		}
	}

	jsUnique, err := json.Marshal(unique)
	if err != nil {
		return nil, err
	}

	jsTotal, err := json.Marshal(total)
	if err != nil {
		return nil, err
	}

	return map[string]interface{}{
		"dates":  dates,
		"unique": string(jsUnique),
		"total":  string(jsTotal),
		"from":   dates[0],
		"to":     dates[len(dates)-1],
	}, nil
}
<commit_msg>adding initial caching of daily metrics from api requests<commit_after>// Package analytics provides the methods to run an analytics reporting system
// for API requests which may be useful to users for measuring access and
// possibly identifying bad actors abusing requests.
package analytics

import (
	"encoding/json"
	"log"
	"net/http"
	"runtime"
	"strings"
	"time"

	"github.com/boltdb/bolt"
)

type apiRequest struct {
	URL        string `json:"url"`
	Method     string `json:"http_method"`
	Origin     string `json:"origin"`
	Proto      string `json:"http_protocol"`
	RemoteAddr string `json:"ip_address"`
	Timestamp  int64  `json:"timestamp"`
	External   bool   `json:"external"`
}

type apiMetric struct {
	Date   string `json:"date"`
	Total  int    `json:"total"`
	Unique int    `json:"unique"`
}

var (
	store       *bolt.DB
	requestChan chan apiRequest
)

// Record queues an apiRequest for metrics
func Record(req *http.Request) {
	external := strings.Contains(req.URL.Path, "/external/")

	r := apiRequest{
		URL:        req.URL.String(),
		Method:     req.Method,
		Origin:     req.Header.Get("Origin"),
		Proto:      req.Proto,
		RemoteAddr: req.RemoteAddr,
		Timestamp:  time.Now().Unix() * 1000,
		External:   external,
	}

	// put r on buffered requestChan to take advantage of batch insertion in DB
	requestChan <- r
}

// Close exports the abillity to close our db file. Should be called with defer
// after call to Init() from the same place.
func Close() {
	err := store.Close()
	if err != nil {
		log.Println(err)
	}
}

// Init creates a db connection, initializes the db with schema and data and
// sets up the queue/batching channel
func Init() {
	var err error
	store, err = bolt.Open("analytics.db", 0666, nil)
	if err != nil {
		log.Fatalln(err)
	}

	err = store.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte("__requests"))
		if err != nil {
			return err
		}

		_, err = tx.CreateBucketIfNotExists([]byte("__metrics"))
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		log.Fatalln("Error idempotently creating requests bucket in analytics.db:", err)
	}

	requestChan = make(chan apiRequest, 1024*64*runtime.NumCPU())

	go serve()

	if err != nil {
		log.Fatalln(err)
	}
}

func serve() {
	// make timer to notify select to batch request insert from requestChan
	// interval: 30 seconds
	apiRequestTimer := time.NewTicker(time.Second * 30)

	// make timer to notify select to remove analytics older than 14 days
	// interval: 1 week
	// TODO: enable analytics backup service to cloud
	pruneThreshold := time.Hour * 24 * 14
	pruneDBTimer := time.NewTicker(pruneThreshold / 2)

	for {
		select {
		case <-apiRequestTimer.C:
			err := batchInsert(requestChan)
			if err != nil {
				log.Println(err)
			}

		case <-pruneDBTimer.C:
			err := batchPrune(pruneThreshold)
			if err != nil {
				log.Println(err)
			}

		case <-time.After(time.Second * 30):
			continue
		}
	}
}

// ChartData returns the map containing decoded javascript needed to chart 2 weeks of data by day
func ChartData() (map[string]interface{}, error) {
	// set thresholds for today and the 13 days preceeding
	times := [14]time.Time{}
	dates := [14]string{}
	now := time.Now()
	today := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC)

	ips := [14]map[string]struct{}{}
	for i := range ips {
		ips[i] = make(map[string]struct{})
	}

	total := [14]int{}
	unique := [14]int{}

	for i := range times {
		// subtract 24 * i hours to make days prior
		dur := time.Duration(24 * i * -1)
		day := today.Add(time.Hour * dur)

		// day threshold is [...n-1-i, n-1, n]
		times[len(times)-1-i] = day
		dates[len(times)-1-i] = day.Format("01/02")
	}

	// get api request analytics and metrics from db
	var requests = []apiRequest{}
	var metrics = [14]apiMetric{}

	err := store.View(func(tx *bolt.Tx) error {
		m := tx.Bucket([]byte("__metrics"))
		b := tx.Bucket([]byte("__requests"))

		err := m.ForEach(func(k, v []byte) error {
			var metric apiMetric
			err := json.Unmarshal(v, &metric)
			if err != nil {
				log.Println("Error decoding api metric json from analytics db:", err)
				return nil
			}

			// if the metric date is in current date range, insert it into
			// metrics array at the position of the date in dates array
			for i := range dates {
				if metric.Date == dates[i] {
					metrics[i] = metric
				}
			}

			return nil
		})
		if err != nil {
			return err
		}

		err = b.ForEach(func(k, v []byte) error {
			var r apiRequest
			err := json.Unmarshal(v, &r)
			if err != nil {
				log.Println("Error decoding api request json from analytics db:", err)
				return nil
			}

			// delete the record in db if it belongs to a day already in metrics,
			// otherwise append it to requests to be analyzed
			d := time.Unix(r.Timestamp/1000, 0).Format("01/02")
			for i := range dates {
				if metrics[i].Date == d {
					err := b.Delete(k)
					if err != nil {
						return err
					}
				} else {
					requests = append(requests, r)
				}
			}

			return nil
		})
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return nil, err
	}

CHECK_REQUEST:
	for i := range requests {
		ts := time.Unix(requests[i].Timestamp/1000, 0)

		for j := range times {
			// if on today, there will be no next iteration to set values for
			// day prior so all valid requests belong to today
			if j == len(times)-1 {
				if ts.After(times[j]) || ts.Equal(times[j]) {
					// do all record keeping
					total[j]++

					if _, ok := ips[j][requests[i].RemoteAddr]; !ok {
						unique[j]++
						ips[j][requests[i].RemoteAddr] = struct{}{}
					}

					continue CHECK_REQUEST
				}
			}

			if ts.Equal(times[j]) {
				// increment total count for current time threshold (day)
				total[j]++

				// if no IP found for current threshold, increment unique and record IP
				if _, ok := ips[j][requests[i].RemoteAddr]; !ok {
					unique[j]++
					ips[j][requests[i].RemoteAddr] = struct{}{}
				}

				continue CHECK_REQUEST
			}

			if ts.Before(times[j]) {
				// check if older than earliest threshold
				if j == 0 {
					continue CHECK_REQUEST
				}

				// increment total count for previous time threshold (day)
				total[j-1]++

				// if no IP found for day prior, increment unique and record IP
				if _, ok := ips[j-1][requests[i].RemoteAddr]; !ok {
					unique[j-1]++
					ips[j-1][requests[i].RemoteAddr] = struct{}{}
				}
			}
		}
	}

	// loop through total and unique to see which dates are accounted for and
	// insert data from metrics array where dates are not
	for i := range metrics {
		if total[i] == 0 {
			total[i] = metrics[i].Total
		}

		if unique[i] == 0 {
			unique[i] = metrics[i].Unique
		}
	}

	// marshal array counts to js arrays for output to chart
	jsUnique, err := json.Marshal(unique)
	if err != nil {
		return nil, err
	}

	jsTotal, err := json.Marshal(total)
	if err != nil {
		return nil, err
	}

	return map[string]interface{}{
		"dates":  dates,
		"unique": string(jsUnique),
		"total":  string(jsTotal),
		"from":   dates[0],
		"to":     dates[len(dates)-1],
	}, nil
}
<|endoftext|>


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

// DNS resolver client: see RFC 1035.

package dns

import (
	"os"
	"net"
	"time"
)

const packErr = "Failed to pack message"
const servErr = "No servers could be reached"

type Resolver struct {
	Servers  []string            // servers to use
	Search   []string            // suffixes to append to local name
	Port     string              // what port to use
	Ndots    int                 // number of dots in name to trigger absolute lookup -- TODO
	Timeout  int                 // seconds before giving up on packet
	Attempts int                 // lost packets before giving up on server
	Rotate   bool                // round robin among servers -- TODO
	Tcp      bool                // use TCP
	Mangle   func([]byte) []byte // mangle the packet
	Rtt      map[string]int64    // Store round trip times
}

// Basic usage pattern for setting up a resolver:
//
//        res := new(Resolver)
//        res.Servers = []string{"127.0.0.1"}           // set the nameserver
//
//        m := new(Msg)                                 // prepare a new message
//        m.MsgHdr.Recursion_desired = true             // header bits
//        m.Question = make([]Question, 1)              // 1 RR in question section
//        m.Question[0] = Question{"miek.nl", TypeSOA, ClassINET}
//        in, err := res.Query(m)                       // Ask the question
//
// Note that message id checking is left to the caller.
func (res *Resolver) Query(q *Msg) (d *Msg, err os.Error) {
	var (
		c    net.Conn
		in   *Msg
		port string
	)
	if len(res.Servers) == 0 {
		return nil, &Error{Error: "No servers defined"}
	}
        if res.Rtt == nil {
                res.Rtt = make(map[string]int64)
        }
	if res.Port == "" {
		port = "53"
	} else {
		port = res.Port
	}

	if q.Id == 0 {
		// No Id sed, set it
		q.Id = Id()
	}
	sending, ok := q.Pack()
	if !ok {
		return nil, &Error{Error: packErr}
	}

	for i := 0; i < len(res.Servers); i++ {
		server := res.Servers[i] + ":" + port
		t := time.Nanoseconds()
		if res.Tcp {
			c, err = net.Dial("tcp", "", server)
		} else {
			c, err = net.Dial("udp", "", server)
		}
		if err != nil {
			continue
		}
		if res.Tcp {
			in, err = exchangeTCP(c, sending, res, true)
		} else {
			in, err = exchangeUDP(c, sending, res, true)
		}
                res.Rtt[server] = time.Nanoseconds() - t

		// Check id in.id != out.id, should be checked in the client!
		c.Close()
		if err != nil {
			continue
		}
		break
	}
	if err != nil {
		return nil, err
	}
	return in, nil
}

// Xfr is used in communicating with *xfr functions.
// This structure is returned on the channel.
type Xfr struct {
	Add bool // true is to be added, otherwise false
	RR
}

// Start an IXFR, q should contain a *Msg with the question
// for an IXFR: "miek.nl" ANY IXFR. RRs that should be added
// have Xfr.Add set to true otherwise it is false.
// Channel m is closed when the IXFR ends.
func (res *Resolver) Ixfr(q *Msg, m chan Xfr) {
	var port string
	var err os.Error
	var in *Msg
	var x Xfr
	if res.Port == "" {
		port = "53"
	} else {
		port = res.Port
	}

	var _ = err // TODO(mg)

	if q.Id == 0 {
		q.Id = Id()
	}

	defer close(m)
	sending, ok := q.Pack()
	if !ok {
		return
	}

Server:
	for i := 0; i < len(res.Servers); i++ {
		server := res.Servers[i] + ":" + port
		c, cerr := net.Dial("tcp", "", server)
		if cerr != nil {
			err = cerr
			continue Server
		}
		first := true
		var serial uint32 // The first serial seen is the current server serial

		defer c.Close()
		for {
			if first {
				in, cerr = exchangeTCP(c, sending, res, true)
			} else {
				in, err = exchangeTCP(c, sending, res, false)
			}

			if cerr != nil {
				// Failed to send, try the next
				err = cerr
				c.Close()
				continue Server
			}
			if in.Id != q.Id {
				return
			}

			if first {
				// A single SOA RR signals "no changes"
				if len(in.Answer) == 1 && checkAxfrSOA(in, true) {
					return
				}

				// But still check if the returned answer is ok
				if !checkAxfrSOA(in, true) {
					c.Close()
					continue Server
				}
				// This serial is important
				serial = in.Answer[0].(*RR_SOA).Serial
				first = !first
			}

			// Now we need to check each message for SOA records, to see what we need to do
			x.Add = true
			if !first {
				for k, r := range in.Answer {
					// If the last record in the IXFR contains the servers' SOA,  we should quit
					if r.Header().Rrtype == TypeSOA {
						switch {
						case r.(*RR_SOA).Serial == serial:
							if k == len(in.Answer)-1 {
								// last rr is SOA with correct serial
								//m <- r dont' send it
								return
							}
							x.Add = true
							if k != 0 {
								// Intermediate SOA
								continue
							}
						case r.(*RR_SOA).Serial != serial:
							x.Add = false
							continue // Don't need to see this SOA
						}
					}
					x.RR = r
					m <- x
				}
			}
			return
		}
		panic("not reached")
		return
	}
	return
}

// Start an AXFR, q should contain a message with the question
// for an AXFR: "miek.nl" ANY AXFR. The closing SOA isn't
// returned over the channel, so the caller will receive
// the zone as-is. Xfr.Add is always true.
// The channel is closed to signal the end of the AXFR.
func (res *Resolver) Axfr(q *Msg, m chan Xfr) {
	var port string
	var err os.Error
	var in *Msg
	if res.Port == "" {
		port = "53"
	} else {
		port = res.Port
	}

	var _ = err // TODO(mg)

	if q.Id == 0 {
		q.Id = Id()
	}

	defer close(m)
	sending, ok := q.Pack()
	if !ok {
		return
	}
Server:
	for i := 0; i < len(res.Servers); i++ {
		server := res.Servers[i] + ":" + port
		c, cerr := net.Dial("tcp", "", server)
		if cerr != nil {
			err = cerr
			continue Server
		}
		first := true
		defer c.Close() // TODO(mg): if not open?
		for {
			if first {
				in, cerr = exchangeTCP(c, sending, res, true)
			} else {
				in, err = exchangeTCP(c, sending, res, false)
			}

			if cerr != nil {
				// Failed to send, try the next
				err = cerr
				c.Close()
				println("AGIAIN")
				continue Server
			}
			if in.Id != q.Id {
				c.Close()
				return
			}

			if first {
				if !checkAxfrSOA(in, true) {
					c.Close()
					continue Server
				}
				first = !first
			}

			if !first {
				if !checkAxfrSOA(in, false) {
					// Soa record not the last one
					sendFromMsg(in, m, false)
					continue
				} else {
					sendFromMsg(in, m, true)
					return
				}
			}
		}
		panic("not reached")
		return
	}
	return
}

// Send a request on the connection and hope for a reply.
// Up to res.Attempts attempts. If send is false, nothing
// is send.
func exchangeUDP(c net.Conn, m []byte, r *Resolver, send bool) (*Msg, os.Error) {
	var timeout int64
	var attempts int
	if r.Mangle != nil {
		m = r.Mangle(m)
	}
	if r.Timeout == 0 {
		timeout = 1
	} else {
		timeout = int64(r.Timeout)
	}
	if r.Attempts == 0 {
		attempts = 1
	} else {
		attempts = r.Attempts
	}
	for a := 0; a < attempts; a++ {
		if send {
			err := sendUDP(m, c)
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return nil, err
			}
		}

		c.SetReadTimeout(timeout * 1e9) // nanoseconds
		buf, err := recvUDP(c)
		if err != nil {
			if e, ok := err.(net.Error); ok && e.Timeout() {
				continue
			}
			return nil, err
		}

		in := new(Msg)
		if !in.Unpack(buf) {
			continue
		}
		return in, nil
	}
	return nil, &Error{Error: servErr}
}

// Up to res.Attempts attempts.
func exchangeTCP(c net.Conn, m []byte, r *Resolver, send bool) (*Msg, os.Error) {
	var timeout int64
	var attempts int
	if r.Mangle != nil {
		m = r.Mangle(m)
	}
	if r.Timeout == 0 {
		timeout = 1
	} else {
		timeout = int64(r.Timeout)
	}
	if r.Attempts == 0 {
		attempts = 1
	} else {
		attempts = r.Attempts
	}

	for a := 0; a < attempts; a++ {
		// only send something when told so
		if send {
			err := sendTCP(m, c)
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return nil, err
			}
		}

		c.SetReadTimeout(timeout * 1e9) // nanoseconds
		// The server replies with two bytes length
		buf, err := recvTCP(c)
		if err != nil {
			if e, ok := err.(net.Error); ok && e.Timeout() {
				continue
			}
			return nil, err
		}
		in := new(Msg)
		if !in.Unpack(buf) {
			continue
		}
		return in, nil
	}
	return nil, &Error{Error: servErr}
}

func sendUDP(m []byte, c net.Conn) os.Error {
	_, err := c.Write(m)
	if err != nil {
		return err
	}
	return nil
}

func recvUDP(c net.Conn) ([]byte, os.Error) {
	m := make([]byte, DefaultMsgSize) // More than enough???
	n, err := c.Read(m)
	if err != nil {
		return nil, err
	}
	m = m[:n]
	return m, nil
}

func sendTCP(m []byte, c net.Conn) os.Error {
	l := make([]byte, 2)
	l[0] = byte(len(m) >> 8)
	l[1] = byte(len(m))
	// First we send the length
	_, err := c.Write(l)
	if err != nil {
		return err
	}
	// And the the message
	_, err = c.Write(m)
	if err != nil {
		return err
	}
	return nil
}

func recvTCP(c net.Conn) ([]byte, os.Error) {
	l := make([]byte, 2) // receiver length
	// The server replies with two bytes length
	_, err := c.Read(l)
	if err != nil {
		return nil, err
	}
	length := uint16(l[0])<<8 | uint16(l[1])
	if length == 0 {
		return nil, &Error{Error: "received nil msg length", Server: c.RemoteAddr().String()}
	}
	m := make([]byte, length)
	n, cerr := c.Read(m)
	if cerr != nil {
		return nil, cerr
	}
	i := n
	for i < int(length) {
		n, err = c.Read(m[i:])
		if err != nil {
			return nil, err
		}
		i += n
	}
	return m, nil
}

// Check if he SOA record exists in the Answer section of
// the packet. If first is true the first RR must be a soa
// if false, the last one should be a SOA
func checkAxfrSOA(in *Msg, first bool) bool {
	if len(in.Answer) > 0 {
		if first {
			return in.Answer[0].Header().Rrtype == TypeSOA
		} else {
			return in.Answer[len(in.Answer)-1].Header().Rrtype == TypeSOA
		}
	}
	return false
}

// Send the answer section to the channel
func sendFromMsg(in *Msg, c chan Xfr, nosoa bool) {
	x := Xfr{Add: true}
	for k, r := range in.Answer {
		if nosoa && k == len(in.Answer)-1 {
			continue
		}
		x.RR = r
		c <- x
	}
}
<commit_msg>Set the map here too<commit_after>// Copyright 2011 Miek Gieben. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// DNS resolver client: see RFC 1035.

package dns

import (
	"os"
	"net"
	"time"
)

const packErr = "Failed to pack message"
const servErr = "No servers could be reached"

type Resolver struct {
	Servers  []string            // servers to use
	Search   []string            // suffixes to append to local name
	Port     string              // what port to use
	Ndots    int                 // number of dots in name to trigger absolute lookup -- TODO
	Timeout  int                 // seconds before giving up on packet
	Attempts int                 // lost packets before giving up on server
	Rotate   bool                // round robin among servers -- TODO
	Tcp      bool                // use TCP
	Mangle   func([]byte) []byte // mangle the packet
	Rtt      map[string]int64    // Store round trip times
}

// Basic usage pattern for setting up a resolver:
//
//        res := new(Resolver)
//        res.Servers = []string{"127.0.0.1"}           // set the nameserver
//
//        m := new(Msg)                                 // prepare a new message
//        m.MsgHdr.Recursion_desired = true             // header bits
//        m.Question = make([]Question, 1)              // 1 RR in question section
//        m.Question[0] = Question{"miek.nl", TypeSOA, ClassINET}
//        in, err := res.Query(m)                       // Ask the question
//
// Note that message id checking is left to the caller.
func (res *Resolver) Query(q *Msg) (d *Msg, err os.Error) {
	var (
		c    net.Conn
		in   *Msg
		port string
	)
	if len(res.Servers) == 0 {
		return nil, &Error{Error: "No servers defined"}
	}
        if res.Rtt == nil {
                res.Rtt = make(map[string]int64)
        }
	if res.Port == "" {
		port = "53"
	} else {
		port = res.Port
	}

	if q.Id == 0 {
		// No Id sed, set it
		q.Id = Id()
	}
	sending, ok := q.Pack()
	if !ok {
		return nil, &Error{Error: packErr}
	}

	for i := 0; i < len(res.Servers); i++ {
		server := res.Servers[i] + ":" + port
		t := time.Nanoseconds()
		if res.Tcp {
			c, err = net.Dial("tcp", "", server)
		} else {
			c, err = net.Dial("udp", "", server)
		}
		if err != nil {
			continue
		}
		if res.Tcp {
			in, err = exchangeTCP(c, sending, res, true)
		} else {
			in, err = exchangeUDP(c, sending, res, true)
		}
                res.Rtt[server] = time.Nanoseconds() - t

		// Check id in.id != out.id, should be checked in the client!
		c.Close()
		if err != nil {
			continue
		}
		break
	}
	if err != nil {
		return nil, err
	}
	return in, nil
}

// Xfr is used in communicating with *xfr functions.
// This structure is returned on the channel.
type Xfr struct {
	Add bool // true is to be added, otherwise false
	RR
}

// Start an IXFR, q should contain a *Msg with the question
// for an IXFR: "miek.nl" ANY IXFR. RRs that should be added
// have Xfr.Add set to true otherwise it is false.
// Channel m is closed when the IXFR ends.
func (res *Resolver) Ixfr(q *Msg, m chan Xfr) {
	var port string
	var err os.Error
	var in *Msg
	var x Xfr
	if res.Port == "" {
		port = "53"
	} else {
		port = res.Port
	}
        if res.Rtt == nil {
                res.Rtt = make(map[string]int64)
        }

	var _ = err // TODO(mg)

	if q.Id == 0 {
		q.Id = Id()
	}

	defer close(m)
	sending, ok := q.Pack()
	if !ok {
		return
	}

Server:
	for i := 0; i < len(res.Servers); i++ {
		server := res.Servers[i] + ":" + port
		c, cerr := net.Dial("tcp", "", server)
		if cerr != nil {
			err = cerr
			continue Server
		}
		first := true
		var serial uint32 // The first serial seen is the current server serial

		defer c.Close()
		for {
			if first {
				in, cerr = exchangeTCP(c, sending, res, true)
			} else {
				in, err = exchangeTCP(c, sending, res, false)
			}

			if cerr != nil {
				// Failed to send, try the next
				err = cerr
				c.Close()
				continue Server
			}
			if in.Id != q.Id {
				return
			}

			if first {
				// A single SOA RR signals "no changes"
				if len(in.Answer) == 1 && checkAxfrSOA(in, true) {
					return
				}

				// But still check if the returned answer is ok
				if !checkAxfrSOA(in, true) {
					c.Close()
					continue Server
				}
				// This serial is important
				serial = in.Answer[0].(*RR_SOA).Serial
				first = !first
			}

			// Now we need to check each message for SOA records, to see what we need to do
			x.Add = true
			if !first {
				for k, r := range in.Answer {
					// If the last record in the IXFR contains the servers' SOA,  we should quit
					if r.Header().Rrtype == TypeSOA {
						switch {
						case r.(*RR_SOA).Serial == serial:
							if k == len(in.Answer)-1 {
								// last rr is SOA with correct serial
								//m <- r dont' send it
								return
							}
							x.Add = true
							if k != 0 {
								// Intermediate SOA
								continue
							}
						case r.(*RR_SOA).Serial != serial:
							x.Add = false
							continue // Don't need to see this SOA
						}
					}
					x.RR = r
					m <- x
				}
			}
			return
		}
		panic("not reached")
		return
	}
	return
}

// Start an AXFR, q should contain a message with the question
// for an AXFR: "miek.nl" ANY AXFR. The closing SOA isn't
// returned over the channel, so the caller will receive
// the zone as-is. Xfr.Add is always true.
// The channel is closed to signal the end of the AXFR.
func (res *Resolver) Axfr(q *Msg, m chan Xfr) {
	var port string
	var err os.Error
	var in *Msg
	if res.Port == "" {
		port = "53"
	} else {
		port = res.Port
	}
        if res.Rtt == nil {
                res.Rtt = make(map[string]int64)
        }

	var _ = err // TODO(mg)

	if q.Id == 0 {
		q.Id = Id()
	}

	defer close(m)
	sending, ok := q.Pack()
	if !ok {
		return
	}
Server:
	for i := 0; i < len(res.Servers); i++ {
		server := res.Servers[i] + ":" + port
		c, cerr := net.Dial("tcp", "", server)
		if cerr != nil {
			err = cerr
			continue Server
		}
		first := true
		defer c.Close() // TODO(mg): if not open?
		for {
			if first {
				in, cerr = exchangeTCP(c, sending, res, true)
			} else {
				in, err = exchangeTCP(c, sending, res, false)
			}

			if cerr != nil {
				// Failed to send, try the next
				err = cerr
				c.Close()
				println("AGIAIN")
				continue Server
			}
			if in.Id != q.Id {
				c.Close()
				return
			}

			if first {
				if !checkAxfrSOA(in, true) {
					c.Close()
					continue Server
				}
				first = !first
			}

			if !first {
				if !checkAxfrSOA(in, false) {
					// Soa record not the last one
					sendFromMsg(in, m, false)
					continue
				} else {
					sendFromMsg(in, m, true)
					return
				}
			}
		}
		panic("not reached")
		return
	}
	return
}

// Send a request on the connection and hope for a reply.
// Up to res.Attempts attempts. If send is false, nothing
// is send.
func exchangeUDP(c net.Conn, m []byte, r *Resolver, send bool) (*Msg, os.Error) {
	var timeout int64
	var attempts int
	if r.Mangle != nil {
		m = r.Mangle(m)
	}
	if r.Timeout == 0 {
		timeout = 1
	} else {
		timeout = int64(r.Timeout)
	}
	if r.Attempts == 0 {
		attempts = 1
	} else {
		attempts = r.Attempts
	}
	for a := 0; a < attempts; a++ {
		if send {
			err := sendUDP(m, c)
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return nil, err
			}
		}

		c.SetReadTimeout(timeout * 1e9) // nanoseconds
		buf, err := recvUDP(c)
		if err != nil {
			if e, ok := err.(net.Error); ok && e.Timeout() {
				continue
			}
			return nil, err
		}

		in := new(Msg)
		if !in.Unpack(buf) {
			continue
		}
		return in, nil
	}
	return nil, &Error{Error: servErr}
}

// Up to res.Attempts attempts.
func exchangeTCP(c net.Conn, m []byte, r *Resolver, send bool) (*Msg, os.Error) {
	var timeout int64
	var attempts int
	if r.Mangle != nil {
		m = r.Mangle(m)
	}
	if r.Timeout == 0 {
		timeout = 1
	} else {
		timeout = int64(r.Timeout)
	}
	if r.Attempts == 0 {
		attempts = 1
	} else {
		attempts = r.Attempts
	}

	for a := 0; a < attempts; a++ {
		// only send something when told so
		if send {
			err := sendTCP(m, c)
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return nil, err
			}
		}

		c.SetReadTimeout(timeout * 1e9) // nanoseconds
		// The server replies with two bytes length
		buf, err := recvTCP(c)
		if err != nil {
			if e, ok := err.(net.Error); ok && e.Timeout() {
				continue
			}
			return nil, err
		}
		in := new(Msg)
		if !in.Unpack(buf) {
			continue
		}
		return in, nil
	}
	return nil, &Error{Error: servErr}
}

func sendUDP(m []byte, c net.Conn) os.Error {
	_, err := c.Write(m)
	if err != nil {
		return err
	}
	return nil
}

func recvUDP(c net.Conn) ([]byte, os.Error) {
	m := make([]byte, DefaultMsgSize) // More than enough???
	n, err := c.Read(m)
	if err != nil {
		return nil, err
	}
	m = m[:n]
	return m, nil
}

func sendTCP(m []byte, c net.Conn) os.Error {
	l := make([]byte, 2)
	l[0] = byte(len(m) >> 8)
	l[1] = byte(len(m))
	// First we send the length
	_, err := c.Write(l)
	if err != nil {
		return err
	}
	// And the the message
	_, err = c.Write(m)
	if err != nil {
		return err
	}
	return nil
}

func recvTCP(c net.Conn) ([]byte, os.Error) {
	l := make([]byte, 2) // receiver length
	// The server replies with two bytes length
	_, err := c.Read(l)
	if err != nil {
		return nil, err
	}
	length := uint16(l[0])<<8 | uint16(l[1])
	if length == 0 {
		return nil, &Error{Error: "received nil msg length", Server: c.RemoteAddr().String()}
	}
	m := make([]byte, length)
	n, cerr := c.Read(m)
	if cerr != nil {
		return nil, cerr
	}
	i := n
	for i < int(length) {
		n, err = c.Read(m[i:])
		if err != nil {
			return nil, err
		}
		i += n
	}
	return m, nil
}

// Check if he SOA record exists in the Answer section of
// the packet. If first is true the first RR must be a soa
// if false, the last one should be a SOA
func checkAxfrSOA(in *Msg, first bool) bool {
	if len(in.Answer) > 0 {
		if first {
			return in.Answer[0].Header().Rrtype == TypeSOA
		} else {
			return in.Answer[len(in.Answer)-1].Header().Rrtype == TypeSOA
		}
	}
	return false
}

// Send the answer section to the channel
func sendFromMsg(in *Msg, c chan Xfr, nosoa bool) {
	x := Xfr{Add: true}
	for k, r := range in.Answer {
		if nosoa && k == len(in.Answer)-1 {
			continue
		}
		x.RR = r
		c <- x
	}
}
<|endoftext|>


<commit_before>package dnsr

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/miekg/dns"
)

// DNS Resolution configuration.
var (
	Timeout             = 2000 * time.Millisecond
	TypicalResponseTime = 100 * time.Millisecond
	MaxRecursion        = 10
	MaxNameservers      = 2
	MaxIPs              = 2
)

// Resolver errors.
var (
	NXDOMAIN = fmt.Errorf("NXDOMAIN")

	ErrMaxRecursion = fmt.Errorf("maximum recursion depth reached: %d", MaxRecursion)
	ErrMaxIPs       = fmt.Errorf("maximum name server IPs queried: %d", MaxIPs)
	ErrNoARecords   = fmt.Errorf("no A records found for name server")
	ErrNoResponse   = fmt.Errorf("no responses received")
	ErrTimeout      = fmt.Errorf("timeout expired") // TODO: Timeouter interface? e.g. func (e) Timeout() bool { return true }
)

// Resolver implements a primitive, non-recursive, caching DNS resolver.
type Resolver struct {
	cache   *cache
	timeout time.Duration
}

// New initializes a Resolver with the specified cache size.
func New(capacity int) *Resolver {
	return NewWithTimeout(capacity, Timeout)
}

// NewWithTimeout initializes a Resolver with the specified cache size and resolution timeout.
func NewWithTimeout(capacity int, timeout time.Duration) *Resolver {
	r := &Resolver{
		cache:   newCache(capacity),
		timeout: timeout,
	}
	return r
}

// Resolve calls ResolveErr to find DNS records of type qtype for the domain qname.
// For nonexistent domains (NXDOMAIN), it will return an empty, non-nil slice.
func (r *Resolver) Resolve(qname, qtype string) RRs {
	rrs, err := r.ResolveErr(qname, qtype)
	if err == NXDOMAIN {
		return emptyRRs
	}
	if err != nil {
		return nil
	}
	return rrs
}

// ResolveErr finds DNS records of type qtype for the domain qname.
// For nonexistent domains, it will return an NXDOMAIN error.
// Specify an empty string in qtype to receive any DNS records found
// (currently A, AAAA, NS, CNAME, SOA, and TXT).
func (r *Resolver) ResolveErr(qname, qtype string) (RRs, error) {
	ctx, cancel := context.WithTimeout(context.Background(), r.timeout)
	defer cancel()
	return r.resolve(ctx, toLowerFQDN(qname), qtype, 0)
}

// ResolveCtx finds DNS records of type qtype for the domain qname using
// the supplied context. Requests may time out earlier if timeout is
// shorter than a deadline set in ctx.
// For nonexistent domains, it will return an NXDOMAIN error.
// Specify an empty string in qtype to receive any DNS records found
// (currently A, AAAA, NS, CNAME, SOA, and TXT).
func (r *Resolver) ResolveCtx(ctx context.Context, qname, qtype string) (RRs, error) {
	ctx, cancel := context.WithTimeout(ctx, r.timeout)
	defer cancel()
	return r.resolve(ctx, toLowerFQDN(qname), qtype, 0)
}

func (r *Resolver) resolve(ctx context.Context, qname, qtype string, depth int) (RRs, error) {
	if depth++; depth > MaxRecursion {
		logMaxRecursion(qname, qtype, depth)
		return nil, ErrMaxRecursion
	}
	rrs, err := r.cacheGet(ctx, qname, qtype)
	if err != nil {
		return nil, err
	}
	if len(rrs) > 0 {
		return rrs, nil
	}
	logResolveStart(qname, qtype, depth)
	start := time.Now()
	rrs, err = r.iterateParents(ctx, qname, qtype, depth)
	logResolveEnd(qname, qtype, rrs, depth, start, err)
	return rrs, err
}

func (r *Resolver) iterateParents(ctx context.Context, qname, qtype string, depth int) (RRs, error) {
	chanRRs := make(chan RRs, MaxNameservers)
	chanErrs := make(chan error, MaxNameservers)
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	for pname, ok := qname, true; ok; pname, ok = parent(pname) {
		// If we’re looking for [foo.com,NS], then move on to the parent ([com,NS])
		if pname == qname && qtype == "NS" {
			continue
		}

		// Only query TLDs against the root nameservers
		if pname == "." && dns.CountLabel(qname) != 1 {
			// fmt.Fprintf(os.Stderr, "Warning: non-TLD query at root: dig +norecurse %s %s\n", qname, qtype)
			return nil, nil
		}

		// Get nameservers
		nrrs, err := r.resolve(ctx, pname, "NS", depth)
		if err == NXDOMAIN || err == ErrTimeout || err == context.DeadlineExceeded {
			return nil, err
		}
		if err != nil {
			continue
		}

		// Check cache for specific queries
		if len(nrrs) > 0 && qtype != "" {
			rrs, err := r.cacheGet(ctx, qname, qtype)
			if err != nil {
				return nil, err
			}
			if len(rrs) > 0 {
				return rrs, nil
			}
		}

		// Query all nameservers in parallel
		count := 0
		for _, nrr := range nrrs {
			if nrr.Type != "NS" {
				continue
			}

			go func(host string) {
				rrs, err := r.exchange(ctx, host, qname, qtype, depth)
				if err != nil {
					chanErrs <- err
				} else {
					chanRRs <- rrs
				}
			}(nrr.Value)

			count++
			if count >= MaxNameservers {
				break
			}
		}

		// Wait for answer, error, or cancellation
		for ; count > 0; count-- {
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case rrs := <-chanRRs:
				for _, nrr := range nrrs {
					if nrr.Name == qname {
						rrs = append(rrs, nrr)
					}
				}
				cancel() // stop any other work here before recursing
				return r.resolveCNAMEs(ctx, qname, qtype, rrs, depth)
			case err = <-chanErrs:
				if err == NXDOMAIN {
					return nil, err
				}
			}
		}

		// NS queries naturally recurse, so stop further iteration
		if qtype == "NS" {
			return nil, err
		}
	}

	return nil, ErrNoResponse
}

func (r *Resolver) exchange(ctx context.Context, host, qname, qtype string, depth int) (RRs, error) {
	dtype := dns.StringToType[qtype]
	if dtype == 0 {
		dtype = dns.TypeA
	}
	qmsg := &dns.Msg{}
	qmsg.SetQuestion(qname, dtype)
	qmsg.MsgHdr.RecursionDesired = false

	// Find each A record for the DNS server
	count := 0
	arrs, err := r.resolve(ctx, host, "A", depth)
	if err != nil {
		return nil, err
	}
	for _, arr := range arrs {
		if arr.Type != "A" { // FIXME: support AAAA records?
			continue
		}

		// Never query more than MaxIPs for any nameserver
		if count++; count > MaxIPs {
			return nil, ErrMaxIPs
		}

		// Synchronously query this DNS server
		start := time.Now()
		timeout := r.timeout // belt and suspenders, since ctx has a deadline from ResolveErr
		if dl, ok := ctx.Deadline(); ok {
			if start.After(dl.Add(-TypicalResponseTime)) { // bail if we can't finish in time (start is too close to deadline)
				return nil, ErrTimeout
			}
			timeout = dl.Sub(start)
		}

		client := &dns.Client{Timeout: timeout} // client must finish within remaining timeout
		rmsg, dur, err := client.Exchange(qmsg, arr.Value+":53")
		select {
		case <-ctx.Done(): // Finished too late
			logCancellation(host, qmsg, rmsg, depth, dur, timeout)
			return nil, ctx.Err()
		default:
			logExchange(host, qmsg, rmsg, depth, dur, timeout, err) // Log hostname instead of IP
		}
		if err != nil {
			continue
		}

		// FIXME: cache NXDOMAIN responses responsibly
		if rmsg.Rcode == dns.RcodeNameError {
			var hasSOA bool
			if qtype == "NS" {
				for _, drr := range rmsg.Ns {
					rr, ok := convertRR(drr)
					if !ok {
						continue
					}
					if rr.Type == "SOA" {
						hasSOA = true
						break
					}
				}
			}
			if !hasSOA {
				r.cache.addNX(qname)
				return nil, NXDOMAIN
			}
		} else if rmsg.Rcode != dns.RcodeSuccess {
			return nil, errors.New(dns.RcodeToString[rmsg.Rcode])
		}

		// Cache records returned
		rrs := r.saveDNSRR(host, qname, append(append(rmsg.Answer, rmsg.Ns...), rmsg.Extra...))

		// Return after first successful network request
		return rrs, nil
	}

	return nil, ErrNoARecords
}

func (r *Resolver) resolveCNAMEs(ctx context.Context, qname, qtype string, crrs RRs, depth int) (RRs, error) {
	var rrs RRs
	for _, crr := range crrs {
		rrs = append(rrs, crr)
		if crr.Type != "CNAME" || crr.Name != qname {
			continue
		}
		logCNAME(crr.String(), depth)
		crrs, _ := r.resolve(ctx, crr.Value, qtype, depth)
		for _, rr := range crrs {
			r.cache.add(qname, rr)
			rrs = append(rrs, crr)
		}
	}
	return rrs, nil
}

// saveDNSRR saves 1 or more DNS records to the resolver cache.
func (r *Resolver) saveDNSRR(host, qname string, drrs []dns.RR) RRs {
	var rrs RRs
	cl := dns.CountLabel(qname)
	for _, drr := range drrs {
		rr, ok := convertRR(drr)
		if !ok {
			continue
		}
		if dns.CountLabel(rr.Name) < cl && dns.CompareDomainName(qname, rr.Name) < 2 {
			// fmt.Fprintf(os.Stderr, "Warning: potential poisoning from %s: %s -> %s\n", host, qname, drr.String())
			continue
		}
		r.cache.add(rr.Name, rr)
		if rr.Name != qname {
			continue
		}
		rrs = append(rrs, rr)
	}
	return rrs
}

// cacheGet returns a randomly ordered slice of DNS records.
func (r *Resolver) cacheGet(ctx context.Context, qname, qtype string) (RRs, error) {
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	default:
	}
	any := r.cache.get(qname)
	if any == nil {
		any = rootCache.get(qname)
	}
	if any == nil {
		return nil, nil
	}
	if len(any) == 0 {
		return nil, NXDOMAIN
	}
	rrs := make(RRs, 0, len(any))
	for _, rr := range any {
		if qtype == "" || rr.Type == qtype {
			rrs = append(rrs, rr)
		}
	}
	if len(rrs) == 0 && (qtype != "" && qtype != "NS") {
		return nil, nil
	}
	return rrs, nil
}
<commit_msg>dnsr: increase MaxNameservers from 2 to 4.<commit_after>package dnsr

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/miekg/dns"
)

// DNS Resolution configuration.
var (
	Timeout             = 2000 * time.Millisecond
	TypicalResponseTime = 100 * time.Millisecond
	MaxRecursion        = 10
	MaxNameservers      = 4
	MaxIPs              = 2
)

// Resolver errors.
var (
	NXDOMAIN = fmt.Errorf("NXDOMAIN")

	ErrMaxRecursion = fmt.Errorf("maximum recursion depth reached: %d", MaxRecursion)
	ErrMaxIPs       = fmt.Errorf("maximum name server IPs queried: %d", MaxIPs)
	ErrNoARecords   = fmt.Errorf("no A records found for name server")
	ErrNoResponse   = fmt.Errorf("no responses received")
	ErrTimeout      = fmt.Errorf("timeout expired") // TODO: Timeouter interface? e.g. func (e) Timeout() bool { return true }
)

// Resolver implements a primitive, non-recursive, caching DNS resolver.
type Resolver struct {
	cache   *cache
	timeout time.Duration
}

// New initializes a Resolver with the specified cache size.
func New(capacity int) *Resolver {
	return NewWithTimeout(capacity, Timeout)
}

// NewWithTimeout initializes a Resolver with the specified cache size and resolution timeout.
func NewWithTimeout(capacity int, timeout time.Duration) *Resolver {
	r := &Resolver{
		cache:   newCache(capacity),
		timeout: timeout,
	}
	return r
}

// Resolve calls ResolveErr to find DNS records of type qtype for the domain qname.
// For nonexistent domains (NXDOMAIN), it will return an empty, non-nil slice.
func (r *Resolver) Resolve(qname, qtype string) RRs {
	rrs, err := r.ResolveErr(qname, qtype)
	if err == NXDOMAIN {
		return emptyRRs
	}
	if err != nil {
		return nil
	}
	return rrs
}

// ResolveErr finds DNS records of type qtype for the domain qname.
// For nonexistent domains, it will return an NXDOMAIN error.
// Specify an empty string in qtype to receive any DNS records found
// (currently A, AAAA, NS, CNAME, SOA, and TXT).
func (r *Resolver) ResolveErr(qname, qtype string) (RRs, error) {
	ctx, cancel := context.WithTimeout(context.Background(), r.timeout)
	defer cancel()
	return r.resolve(ctx, toLowerFQDN(qname), qtype, 0)
}

// ResolveCtx finds DNS records of type qtype for the domain qname using
// the supplied context. Requests may time out earlier if timeout is
// shorter than a deadline set in ctx.
// For nonexistent domains, it will return an NXDOMAIN error.
// Specify an empty string in qtype to receive any DNS records found
// (currently A, AAAA, NS, CNAME, SOA, and TXT).
func (r *Resolver) ResolveCtx(ctx context.Context, qname, qtype string) (RRs, error) {
	ctx, cancel := context.WithTimeout(ctx, r.timeout)
	defer cancel()
	return r.resolve(ctx, toLowerFQDN(qname), qtype, 0)
}

func (r *Resolver) resolve(ctx context.Context, qname, qtype string, depth int) (RRs, error) {
	if depth++; depth > MaxRecursion {
		logMaxRecursion(qname, qtype, depth)
		return nil, ErrMaxRecursion
	}
	rrs, err := r.cacheGet(ctx, qname, qtype)
	if err != nil {
		return nil, err
	}
	if len(rrs) > 0 {
		return rrs, nil
	}
	logResolveStart(qname, qtype, depth)
	start := time.Now()
	rrs, err = r.iterateParents(ctx, qname, qtype, depth)
	logResolveEnd(qname, qtype, rrs, depth, start, err)
	return rrs, err
}

func (r *Resolver) iterateParents(ctx context.Context, qname, qtype string, depth int) (RRs, error) {
	chanRRs := make(chan RRs, MaxNameservers)
	chanErrs := make(chan error, MaxNameservers)
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	for pname, ok := qname, true; ok; pname, ok = parent(pname) {
		// If we’re looking for [foo.com,NS], then move on to the parent ([com,NS])
		if pname == qname && qtype == "NS" {
			continue
		}

		// Only query TLDs against the root nameservers
		if pname == "." && dns.CountLabel(qname) != 1 {
			// fmt.Fprintf(os.Stderr, "Warning: non-TLD query at root: dig +norecurse %s %s\n", qname, qtype)
			return nil, nil
		}

		// Get nameservers
		nrrs, err := r.resolve(ctx, pname, "NS", depth)
		if err == NXDOMAIN || err == ErrTimeout || err == context.DeadlineExceeded {
			return nil, err
		}
		if err != nil {
			continue
		}

		// Check cache for specific queries
		if len(nrrs) > 0 && qtype != "" {
			rrs, err := r.cacheGet(ctx, qname, qtype)
			if err != nil {
				return nil, err
			}
			if len(rrs) > 0 {
				return rrs, nil
			}
		}

		// Query all nameservers in parallel
		count := 0
		for _, nrr := range nrrs {
			if nrr.Type != "NS" {
				continue
			}

			go func(host string) {
				rrs, err := r.exchange(ctx, host, qname, qtype, depth)
				if err != nil {
					chanErrs <- err
				} else {
					chanRRs <- rrs
				}
			}(nrr.Value)

			count++
			if count >= MaxNameservers {
				break
			}
		}

		// Wait for answer, error, or cancellation
		for ; count > 0; count-- {
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case rrs := <-chanRRs:
				for _, nrr := range nrrs {
					if nrr.Name == qname {
						rrs = append(rrs, nrr)
					}
				}
				cancel() // stop any other work here before recursing
				return r.resolveCNAMEs(ctx, qname, qtype, rrs, depth)
			case err = <-chanErrs:
				if err == NXDOMAIN {
					return nil, err
				}
			}
		}

		// NS queries naturally recurse, so stop further iteration
		if qtype == "NS" {
			return nil, err
		}
	}

	return nil, ErrNoResponse
}

func (r *Resolver) exchange(ctx context.Context, host, qname, qtype string, depth int) (RRs, error) {
	dtype := dns.StringToType[qtype]
	if dtype == 0 {
		dtype = dns.TypeA
	}
	qmsg := &dns.Msg{}
	qmsg.SetQuestion(qname, dtype)
	qmsg.MsgHdr.RecursionDesired = false

	// Find each A record for the DNS server
	count := 0
	arrs, err := r.resolve(ctx, host, "A", depth)
	if err != nil {
		return nil, err
	}
	for _, arr := range arrs {
		if arr.Type != "A" { // FIXME: support AAAA records?
			continue
		}

		// Never query more than MaxIPs for any nameserver
		if count++; count > MaxIPs {
			return nil, ErrMaxIPs
		}

		// Synchronously query this DNS server
		start := time.Now()
		timeout := r.timeout // belt and suspenders, since ctx has a deadline from ResolveErr
		if dl, ok := ctx.Deadline(); ok {
			if start.After(dl.Add(-TypicalResponseTime)) { // bail if we can't finish in time (start is too close to deadline)
				return nil, ErrTimeout
			}
			timeout = dl.Sub(start)
		}

		client := &dns.Client{Timeout: timeout} // client must finish within remaining timeout
		rmsg, dur, err := client.Exchange(qmsg, arr.Value+":53")
		select {
		case <-ctx.Done(): // Finished too late
			logCancellation(host, qmsg, rmsg, depth, dur, timeout)
			return nil, ctx.Err()
		default:
			logExchange(host, qmsg, rmsg, depth, dur, timeout, err) // Log hostname instead of IP
		}
		if err != nil {
			continue
		}

		// FIXME: cache NXDOMAIN responses responsibly
		if rmsg.Rcode == dns.RcodeNameError {
			var hasSOA bool
			if qtype == "NS" {
				for _, drr := range rmsg.Ns {
					rr, ok := convertRR(drr)
					if !ok {
						continue
					}
					if rr.Type == "SOA" {
						hasSOA = true
						break
					}
				}
			}
			if !hasSOA {
				r.cache.addNX(qname)
				return nil, NXDOMAIN
			}
		} else if rmsg.Rcode != dns.RcodeSuccess {
			return nil, errors.New(dns.RcodeToString[rmsg.Rcode])
		}

		// Cache records returned
		rrs := r.saveDNSRR(host, qname, append(append(rmsg.Answer, rmsg.Ns...), rmsg.Extra...))

		// Return after first successful network request
		return rrs, nil
	}

	return nil, ErrNoARecords
}

func (r *Resolver) resolveCNAMEs(ctx context.Context, qname, qtype string, crrs RRs, depth int) (RRs, error) {
	var rrs RRs
	for _, crr := range crrs {
		rrs = append(rrs, crr)
		if crr.Type != "CNAME" || crr.Name != qname {
			continue
		}
		logCNAME(crr.String(), depth)
		crrs, _ := r.resolve(ctx, crr.Value, qtype, depth)
		for _, rr := range crrs {
			r.cache.add(qname, rr)
			rrs = append(rrs, crr)
		}
	}
	return rrs, nil
}

// saveDNSRR saves 1 or more DNS records to the resolver cache.
func (r *Resolver) saveDNSRR(host, qname string, drrs []dns.RR) RRs {
	var rrs RRs
	cl := dns.CountLabel(qname)
	for _, drr := range drrs {
		rr, ok := convertRR(drr)
		if !ok {
			continue
		}
		if dns.CountLabel(rr.Name) < cl && dns.CompareDomainName(qname, rr.Name) < 2 {
			// fmt.Fprintf(os.Stderr, "Warning: potential poisoning from %s: %s -> %s\n", host, qname, drr.String())
			continue
		}
		r.cache.add(rr.Name, rr)
		if rr.Name != qname {
			continue
		}
		rrs = append(rrs, rr)
	}
	return rrs
}

// cacheGet returns a randomly ordered slice of DNS records.
func (r *Resolver) cacheGet(ctx context.Context, qname, qtype string) (RRs, error) {
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	default:
	}
	any := r.cache.get(qname)
	if any == nil {
		any = rootCache.get(qname)
	}
	if any == nil {
		return nil, nil
	}
	if len(any) == 0 {
		return nil, NXDOMAIN
	}
	rrs := make(RRs, 0, len(any))
	for _, rr := range any {
		if qtype == "" || rr.Type == qtype {
			rrs = append(rrs, rr)
		}
	}
	if len(rrs) == 0 && (qtype != "" && qtype != "NS") {
		return nil, nil
	}
	return rrs, nil
}
<|endoftext|>


<commit_before>package erede

import (
	"encoding/xml"
)

const (
	RStatSuccess                       = 1
	RStatSocketWriteError              = 2
	RStatTimeout                       = 3
	RStatEditError                     = 5
	RStatCommsError                    = 6
	RStatUnauthorized                  = 7
	RStatCurrencyError                 = 9
	RStatAuthError                     = 10
	RStatInvalidAuthCode               = 12
	RStatTypeFieldMissing              = 13
	RStatDBServerError                 = 14
	RStatInvalidType                   = 15
	RStatCannotFulfillTransaction      = 19
	RStatDuplicateTransactionReference = 20
	RStatInvalidCardType               = 21
	RStatInvalidReference              = 22
	RStatExpiryDateInvalid             = 23
	RStatCardExpired                   = 24
	RStatCardNumberInvalid             = 25
	RStatCardNumberWrongLength         = 26
	RStatIssueNumberError              = 27
	RStatStartDateError                = 28
	RStatCardNotValidYet               = 29
	RStatStartDateAfterExpiryDate      = 30
	//TODO: fill more errors
	RStatCurrencyNotSupportedByCard = 59
	RStatInvalidXML                 = 60
	//TODO: fill more errors
	RStatPaymentGatewayBusy = 440
	//TODO: fill more errors
	RStatInvalidTransactionType    = 473
	RStatInvalidValueForMerchantID = 480
)

const (
	RRejectedServiceUnauthorized = 51
	//TODO: fill more errors
	RRejectedInvalidVendor   = 57
	RRejectedUnauthorized    = 58
	RRejectedInvalidPassword = 65
	//TODO: fill more errors
)

type TransactionResponse struct {
	XMLName         xml.Name           `xml:"Response"`
	Card            TrRespCard         `xml:"Card"`
	CardTxn         TrRespCardTxn      `xml:"CardTxn"`
	ThreeDSecure    TrRespThreeDSecure `xml:"ThreeDSecure"`
	Acquirer        string             `xml:"acquirer"`
	AuthHostRef     int                `xml:"auth_host_reference"`
	GatewayRef      string             `xml:"gateway_reference"`
	ExtendedRespMsg string             `xml:"extended_response_message"`
	ExtendedStatus  string             `xml:"extended_status"`
	MerchantRef     string             `xml:"merchant_reference"`
	MID             string             `xml:"mid"`
	Mode            string             `xml:"mode"`
	Reason          string             `xml:"reason"`
	Status          int                `xml:"status"`
	Time            int64              `xml:"time"`
}

//TODO: CHECK xid, aav, caavAlgorithm, eci
type TrRespThreeDSecure struct {
	AcsURL        string `xml:"acs_url"`
	PareqMessage  string `xml:"pareq_message"`
	XID           string `xml:"xid"`
	AAV           string `xml:"aav"`
	CAVVAlgorithm string `xml:"cavvAlgorithm"`
	ECI           string `xml:"eci"`
}

type TrRespCard struct {
	AccType string `xml:"card_account_type"`
}

type TrRespCardTxn struct {
	AccType    string              `xml:"card_account_type"`
	Cv2Avs     TrRespCardTxnCv2AVS `xml:"Cv2Avs"`
	AuthCode   string              `xml:"authcode"`
	CardScheme string              `xml:"card_scheme"`
	Country    string              `xml:"country"`
	Issuer     string              `xml:"issuer"`
}

type TrRespCardTxnCv2AVS struct {
	Status string `xml:"cv2avs_status"`
	Policy int    `xml:"policy"`
}

type TransactionResponse2 struct {
	XMLName         xml.Name           `xml:"Response"`
	QueryTxnResult  RespQueryTxnResult `xml:"QueryTxnResult"`
	ExtendedRespMsg string             `xml:"extended_response_message"`
	ExtendedStatus  string             `xml:"extended_status"`
	Mode            string             `xml:"mode"`
	Reason          string             `xml:"reason"`
	Status          int                `xml:"status"`
	Time            int64              `xml:"time"`
}

type RespQueryTxnResult struct {
	Card                 TrRespCard2 `xml:"Card"`
	Acquirer             string      `xml:"acquirer"`
	AuthHostRef          int         `xml:"auth_host_reference"`
	AuthCode             string      `xml:"authcode"`
	GatewayRef           string      `xml:"gateway_reference"`
	Environment          string      `xml:"environment"`
	FulfillDate          string      `xml:"fulfill_date"`
	FulfillTimestamp     int64       `xml:"fulfill_timestamp"`
	MerchantRef          int         `xml:"merchant_reference"`
	Reason               string      `xml:"reason"`
	Sent                 string      `xml:"sent"`
	Status               int         `xml:"status"`
	TransactionDate      string      `xml:"transaction_date"`
	TransactionTimestamp int64       `xml:"transaction_timestamp"`
}

type TrRespCard2 struct {
	Category   string `xml:"card_category"`
	Country    string `xml:"country"`
	ExpiryDate string `xml:"expirydate"`
	Issuer     string `xml:"issuer"`
	PAN        string `xml:"pan"`
	Scheme     string `xml:"scheme"`
}

func GetRespRejectionDescription(code int) string {
	switch code {
	case 51:
		return "Produto ou serviço não habilitado para o estabelecimento. Entre em contato com a Rede."
	case 53:
		return "Transação não permitida para o emissor. Entre em contato com a Rede."
	case 56:
		return "Erro nos dados informados. Tente novamente."
	case 57:
		return "Estabelecimento inválido."
	case 58:
		return "Transação não autorizada. Contate o emissor."
	case 65:
		return "Senha inválida. Tente novamente."
	case 69:
		return "Transação não permitida para este produto ou serviço."
	case 72:
		return "Contate o emissor."
	case 74:
		return "Falha na comunicação. Tente novamente."
	case 79:
		return "Cartão expirado. Transação não pode ser resubmetida. Contate o emissor."
	case 80:
		return "Transação não autorizada. Contate o emissor. (Saldo Insuficiente)"
	case 81:
		return "Produto ou serviço não habilitado para o emissor (AVS)."
	case 82:
		return "Transação não autorizada para cartão de débito."
	case 83:
		return "Transação não autorizada. Problemas com cartão. Contate o emissor."
	case 84:
		return "Transação não autorizada. Transação não pode ser resubmetida. Contate o emissor."
	}
	return "ERRO!"
}

func GetGenRespDescription(code int) string {
	switch code {
	case 1:
		return "Sucesso."
	case 2:
		return "A comunicação foi interrompida"
	case 3:
		return "Ocorreu um timeout enquanto os detalhes da transação eram lidos"
	case 5:
		return "Um campo foi especificado duas vezes. Foram enviados dados excessivos ou inválidos, um fulfill de pré-autorização falhou ou um campo foi omitido. O argumento oferecerá uma melhor indicação do que exatamente deu errado"
	case 6:
		return "Erro no link de comunicação; reenvie"
	case 9:
		return "A moeda especificada não existe"
	case 10:
		return "O vTID ou senha são incorretos"
	case 12:
		return "O código de autorização fornecido é inválido"
	case 13:
		return "Não foi inserido um tipo de transação"
	case 14:
		return "Os detalhes da transação não foram enviados ao nosso banco de dados"
	case 15:
		return "Foi especificado um tipo de transação inválido"
	case 19:
		return "Houve uma tentativa de fulfill de uma transação que não pode ser confirmada ou que já foi confirmada"
	case 20:
		return "Já foi enviada uma transação bem-sucedida que utiliza este vTID e número de referência"
	case 21:
		return "Este terminal não aceita transações para este tipo de cartão"
	case 22:
		return "Os números de referência devem ter 16 dígitos para transações de fulfill, ou de 6 a 30 dígitos para todas as outras"
	case 23:
		return "Expiry date do cartão inválido."
	case 24:
		return "A data de validade fornecida é anterior à data atual"
	case 25, 26:
		return "Número do cartão inválido"
	}
	return ""
}
<commit_msg>added some 3-DS status codes<commit_after>package erede

import (
	"encoding/xml"
)

const (
	RStatSuccess                       = 1
	RStatSocketWriteError              = 2
	RStatTimeout                       = 3
	RStatEditError                     = 5
	RStatCommsError                    = 6
	RStatUnauthorized                  = 7
	RStatCurrencyError                 = 9
	RStatAuthError                     = 10
	RStatInvalidAuthCode               = 12
	RStatTypeFieldMissing              = 13
	RStatDBServerError                 = 14
	RStatInvalidType                   = 15
	RStatCannotFulfillTransaction      = 19
	RStatDuplicateTransactionReference = 20
	RStatInvalidCardType               = 21
	RStatInvalidReference              = 22
	RStatExpiryDateInvalid             = 23
	RStatCardExpired                   = 24
	RStatCardNumberInvalid             = 25
	RStatCardNumberWrongLength         = 26
	RStatIssueNumberError              = 27
	RStatStartDateError                = 28
	RStatCardNotValidYet               = 29
	RStatStartDateAfterExpiryDate      = 30
	//TODO: fill more errors
	RStatCurrencyNotSupportedByCard = 59
	RStatInvalidXML                 = 60
	//TODO: fill more errors
	// 3DS
	RStat3DSPayerVerificationRequired = 150
	RStat3DSInvalidTransactionType    = 151
	RStat3DSManualAuthNotSupported    = 152
	RStat3DSVerifyElmtMissing         = 153
	RStat3DSInvalidVerifyValue        = 154
	RStat3DSFieldMissing              = 155
	RStat3DSInvalidBrowserDeviceCateg = 156
	RStat3DSMerchantNotEnabled        = 157
	RStat3DSSchemeNotSupported        = 158
	RStat3DSVerificationFailed        = 159
	RStat3DSInvalidIssuerResponse     = 160
	RStat3DSAuthFailedCallCentre      = 161
	RStat3DSCardNotEnrolled           = 162
	//TODO: fill more errors
	RStatPaymentGatewayBusy = 440
	//TODO: fill more errors
	RStat3DSRequired               = 471
	RStatInvalidTransactionType    = 473
	RStatInvalidValueForMerchantID = 480
)

const (
	RRejectedServiceUnauthorized = 51
	//TODO: fill more errors
	RRejectedInvalidVendor   = 57
	RRejectedUnauthorized    = 58
	RRejectedInvalidPassword = 65
	//TODO: fill more errors
)

type TransactionResponse struct {
	XMLName         xml.Name           `xml:"Response"`
	Card            TrRespCard         `xml:"Card"`
	CardTxn         TrRespCardTxn      `xml:"CardTxn"`
	ThreeDSecure    TrRespThreeDSecure `xml:"ThreeDSecure"`
	Acquirer        string             `xml:"acquirer"`
	AuthHostRef     int                `xml:"auth_host_reference"`
	GatewayRef      string             `xml:"gateway_reference"`
	ExtendedRespMsg string             `xml:"extended_response_message"`
	ExtendedStatus  string             `xml:"extended_status"`
	MerchantRef     string             `xml:"merchant_reference"`
	MID             string             `xml:"mid"`
	Mode            string             `xml:"mode"`
	Reason          string             `xml:"reason"`
	Status          int                `xml:"status"`
	Time            int64              `xml:"time"`
}

//TODO: CHECK xid, aav, caavAlgorithm, eci
type TrRespThreeDSecure struct {
	AcsURL        string `xml:"acs_url"`
	PareqMessage  string `xml:"pareq_message"`
	XID           string `xml:"xid"`
	AAV           string `xml:"aav"`
	CAVVAlgorithm string `xml:"cavvAlgorithm"`
	ECI           string `xml:"eci"`
}

type TrRespCard struct {
	AccType string `xml:"card_account_type"`
}

type TrRespCardTxn struct {
	AccType    string              `xml:"card_account_type"`
	Cv2Avs     TrRespCardTxnCv2AVS `xml:"Cv2Avs"`
	AuthCode   string              `xml:"authcode"`
	CardScheme string              `xml:"card_scheme"`
	Country    string              `xml:"country"`
	Issuer     string              `xml:"issuer"`
}

type TrRespCardTxnCv2AVS struct {
	Status string `xml:"cv2avs_status"`
	Policy int    `xml:"policy"`
}

type TransactionResponse2 struct {
	XMLName         xml.Name           `xml:"Response"`
	QueryTxnResult  RespQueryTxnResult `xml:"QueryTxnResult"`
	ExtendedRespMsg string             `xml:"extended_response_message"`
	ExtendedStatus  string             `xml:"extended_status"`
	Mode            string             `xml:"mode"`
	Reason          string             `xml:"reason"`
	Status          int                `xml:"status"`
	Time            int64              `xml:"time"`
}

type RespQueryTxnResult struct {
	Card                 TrRespCard2 `xml:"Card"`
	Acquirer             string      `xml:"acquirer"`
	AuthHostRef          int         `xml:"auth_host_reference"`
	AuthCode             string      `xml:"authcode"`
	GatewayRef           string      `xml:"gateway_reference"`
	Environment          string      `xml:"environment"`
	FulfillDate          string      `xml:"fulfill_date"`
	FulfillTimestamp     int64       `xml:"fulfill_timestamp"`
	MerchantRef          int         `xml:"merchant_reference"`
	Reason               string      `xml:"reason"`
	Sent                 string      `xml:"sent"`
	Status               int         `xml:"status"`
	TransactionDate      string      `xml:"transaction_date"`
	TransactionTimestamp int64       `xml:"transaction_timestamp"`
}

type TrRespCard2 struct {
	Category   string `xml:"card_category"`
	Country    string `xml:"country"`
	ExpiryDate string `xml:"expirydate"`
	Issuer     string `xml:"issuer"`
	PAN        string `xml:"pan"`
	Scheme     string `xml:"scheme"`
}

func GetRespRejectionDescription(code int) string {
	switch code {
	case 51:
		return "Produto ou serviço não habilitado para o estabelecimento. Entre em contato com a Rede."
	case 53:
		return "Transação não permitida para o emissor. Entre em contato com a Rede."
	case 56:
		return "Erro nos dados informados. Tente novamente."
	case 57:
		return "Estabelecimento inválido."
	case 58:
		return "Transação não autorizada. Contate o emissor."
	case 65:
		return "Senha inválida. Tente novamente."
	case 69:
		return "Transação não permitida para este produto ou serviço."
	case 72:
		return "Contate o emissor."
	case 74:
		return "Falha na comunicação. Tente novamente."
	case 79:
		return "Cartão expirado. Transação não pode ser resubmetida. Contate o emissor."
	case 80:
		return "Transação não autorizada. Contate o emissor. (Saldo Insuficiente)"
	case 81:
		return "Produto ou serviço não habilitado para o emissor (AVS)."
	case 82:
		return "Transação não autorizada para cartão de débito."
	case 83:
		return "Transação não autorizada. Problemas com cartão. Contate o emissor."
	case 84:
		return "Transação não autorizada. Transação não pode ser resubmetida. Contate o emissor."
	}
	return "ERRO!"
}

func GetGenRespDescription(code int) string {
	switch code {
	case 1:
		return "Sucesso."
	case 2:
		return "A comunicação foi interrompida"
	case 3:
		return "Ocorreu um timeout enquanto os detalhes da transação eram lidos"
	case 5:
		return "Um campo foi especificado duas vezes. Foram enviados dados excessivos ou inválidos, um fulfill de pré-autorização falhou ou um campo foi omitido. O argumento oferecerá uma melhor indicação do que exatamente deu errado"
	case 6:
		return "Erro no link de comunicação; reenvie"
	case 9:
		return "A moeda especificada não existe"
	case 10:
		return "O vTID ou senha são incorretos"
	case 12:
		return "O código de autorização fornecido é inválido"
	case 13:
		return "Não foi inserido um tipo de transação"
	case 14:
		return "Os detalhes da transação não foram enviados ao nosso banco de dados"
	case 15:
		return "Foi especificado um tipo de transação inválido"
	case 19:
		return "Houve uma tentativa de fulfill de uma transação que não pode ser confirmada ou que já foi confirmada"
	case 20:
		return "Já foi enviada uma transação bem-sucedida que utiliza este vTID e número de referência"
	case 21:
		return "Este terminal não aceita transações para este tipo de cartão"
	case 22:
		return "Os números de referência devem ter 16 dígitos para transações de fulfill, ou de 6 a 30 dígitos para todas as outras"
	case 23:
		return "Expiry date do cartão inválido."
	case 24:
		return "A data de validade fornecida é anterior à data atual"
	case 25, 26:
		return "Número do cartão inválido"
	}
	return ""
}
<|endoftext|>


<commit_before>package osin

import (
	"net/http"
	"time"
)

// AccessRequestType is the type for OAuth param `grant_type`
type AccessRequestType string

const (
	AUTHORIZATION_CODE AccessRequestType = "authorization_code"
	REFRESH_TOKEN                        = "refresh_token"
	PASSWORD                             = "password"
	CLIENT_CREDENTIALS                   = "client_credentials"
	IMPLICIT                             = "__implicit"
)

// AccessRequest is a request for access tokens
type AccessRequest struct {
	Type          AccessRequestType
	Code          string
	Client        *Client
	AuthorizeData *AuthorizeData
	AccessData    *AccessData
	RedirectUri   string
	Scope         string
	Username      string
	Password      string

	// Set if request is authorized
	Authorized bool

	// Token expiration in seconds. Change if different from default
	Expiration int32

	// Set if a refresh token should be generated
	GenerateRefresh bool

	// Data to be passed to storage. Not used by the library.
	UserData interface{}
}

// AccessData represents an access grant (tokens, expiration, client, etc)
type AccessData struct {
	// Client information
	Client *Client

	// Authorize data, for authorization code
	AuthorizeData *AuthorizeData

	// Previous access data, for refresh token
	AccessData *AccessData

	// Access token
	AccessToken string

	// Refresh Token. Can be blank
	RefreshToken string

	// Token expiration in seconds
	ExpiresIn int32

	// Requested scope
	Scope string

	// Redirect Uri from request
	RedirectUri string

	// Date created
	CreatedAt time.Time

	// Data to be passed to storage. Not used by the library.
	UserData interface{}
}

// IsExpired returns true if access expired
func (d *AccessData) IsExpired() bool {
	return d.CreatedAt.Add(time.Duration(d.ExpiresIn) * time.Second).Before(time.Now())
}

// ExpireAt returns the expiration date
func (d *AccessData) ExpireAt() time.Time {
	return d.CreatedAt.Add(time.Duration(d.ExpiresIn) * time.Second)
}

// AccessTokenGen generates access tokens
type AccessTokenGen interface {
	GenerateAccessToken(data *AccessData, generaterefresh bool) (accesstoken string, refreshtoken string, err error)
}

// HandleAccessRequest is the http.HandlerFunc for handling access token requests
func (s *Server) HandleAccessRequest(w *Response, r *http.Request) *AccessRequest {
	// Only allow GET or POST
	if r.Method == "GET" {
		if !s.Config.AllowGetAccessRequest {
			w.SetError(E_INVALID_REQUEST, "")
			return nil
		}
	} else if r.Method != "POST" {
		w.SetError(E_INVALID_REQUEST, "")
		return nil
	}

	err := r.ParseForm()
	if err != nil {
		w.SetError(E_INVALID_REQUEST, "")
		return nil
	}

	grantType := AccessRequestType(r.Form.Get("grant_type"))
	if s.Config.AllowedAccessTypes.Exists(grantType) {
		switch grantType {
		case AUTHORIZATION_CODE:
			return s.handleAuthorizationCodeRequest(w, r)
		case REFRESH_TOKEN:
			return s.handleRefreshTokenRequest(w, r)
		case PASSWORD:
			return s.handlePasswordRequest(w, r)
		case CLIENT_CREDENTIALS:
			return s.handleClientCredentialsRequest(w, r)
		}
	}

	w.SetError(E_UNSUPPORTED_GRANT_TYPE, "")
	return nil
}

func (s *Server) handleAuthorizationCodeRequest(w *Response, r *http.Request) *AccessRequest {
	// get client authentication
	auth := getClientAuth(w, r, s.Config.AllowClientSecretInParams)
	if auth == nil {
		return nil
	}

	// generate access token
	ret := &AccessRequest{
		Type:            AUTHORIZATION_CODE,
		Code:            r.Form.Get("code"),
		RedirectUri:     r.Form.Get("redirect_uri"),
		GenerateRefresh: true,
		Expiration:      s.Config.AccessExpiration,
	}

	// "code" is required
	if ret.Code == "" {
		w.SetError(E_INVALID_GRANT, "")
		return nil
	}

	// must have a valid client
	if ret.Client = getClient(auth, s.Storage, w); ret.Client == nil {
		return nil
	}

	// must be a valid authorization code
	var err error
	ret.AuthorizeData, err = s.Storage.LoadAuthorize(ret.Code)
	if err != nil {
		w.SetError(E_INVALID_GRANT, "")
		w.InternalError = err
		return nil
	}
	if ret.AuthorizeData.Client == nil {
		w.SetError(E_UNAUTHORIZED_CLIENT, "")
		return nil
	}
	if ret.AuthorizeData.Client.RedirectUri == "" {
		w.SetError(E_UNAUTHORIZED_CLIENT, "")
		return nil
	}
	if ret.AuthorizeData.IsExpired() {
		w.SetError(E_INVALID_GRANT, "")
		return nil
	}

	// code must be from the client
	if ret.AuthorizeData.Client.Id != ret.Client.Id {
		w.SetError(E_INVALID_GRANT, "")
		return nil
	}

	// check redirect uri
	if ret.RedirectUri == "" {
		ret.RedirectUri = ret.Client.RedirectUri
	}
	if err = ValidateUri(ret.Client.RedirectUri, ret.RedirectUri); err != nil {
		w.SetError(E_INVALID_REQUEST, "")
		w.InternalError = err
		return nil
	}
	if ret.AuthorizeData.RedirectUri != ret.RedirectUri {
		w.SetError(E_INVALID_REQUEST, "")
		return nil
	}

	// set rest of data
	ret.Scope = ret.AuthorizeData.Scope

	return ret
}

func (s *Server) handleRefreshTokenRequest(w *Response, r *http.Request) *AccessRequest {
	// get client authentication
	auth := getClientAuth(w, r, s.Config.AllowClientSecretInParams)
	if auth == nil {
		return nil
	}

	// generate access token
	ret := &AccessRequest{
		Type:            REFRESH_TOKEN,
		Code:            r.Form.Get("refresh_token"),
		Scope:           r.Form.Get("scope"),
		GenerateRefresh: true,
		Expiration:      s.Config.AccessExpiration,
	}

	// "refresh_token" is required
	if ret.Code == "" {
		w.SetError(E_INVALID_GRANT, "")
		return nil
	}

	// must have a valid client
	if ret.Client = getClient(auth, s.Storage, w); ret.Client == nil {
		return nil
	}

	// must be a valid refresh code
	var err error
	ret.AccessData, err = s.Storage.LoadRefresh(ret.Code)
	if err != nil {
		w.SetError(E_INVALID_GRANT, "")
		w.InternalError = err
		return nil
	}
	if ret.AccessData.Client == nil {
		w.SetError(E_UNAUTHORIZED_CLIENT, "")
		return nil
	}
	if ret.AccessData.Client.RedirectUri == "" {
		w.SetError(E_UNAUTHORIZED_CLIENT, "")
		return nil
	}

	// client must be the same as the previous token
	if ret.AccessData.Client.Id != ret.Client.Id {
		w.SetError(E_INVALID_CLIENT, "")
		return nil

	}

	// set rest of data
	ret.RedirectUri = ret.AccessData.RedirectUri
	ret.UserData = ret.AccessData.UserData
	if ret.Scope == "" {
		ret.Scope = ret.AccessData.Scope
	}

	return ret
}

func (s *Server) handlePasswordRequest(w *Response, r *http.Request) *AccessRequest {
	// get client authentication
	auth := getClientAuth(w, r, s.Config.AllowClientSecretInParams)
	if auth == nil {
		return nil
	}

	// generate access token
	ret := &AccessRequest{
		Type:            PASSWORD,
		Username:        r.Form.Get("username"),
		Password:        r.Form.Get("password"),
		Scope:           r.Form.Get("scope"),
		GenerateRefresh: true,
		Expiration:      s.Config.AccessExpiration,
	}

	// "username" and "password" is required
	if ret.Username == "" || ret.Password == "" {
		w.SetError(E_INVALID_GRANT, "")
		return nil
	}

	// must have a valid client
	if ret.Client = getClient(auth, s.Storage, w); ret.Client == nil {
		return nil
	}

	// set redirect uri
	ret.RedirectUri = ret.Client.RedirectUri

	return ret
}

func (s *Server) handleClientCredentialsRequest(w *Response, r *http.Request) *AccessRequest {
	// get client authentication
	auth := getClientAuth(w, r, s.Config.AllowClientSecretInParams)
	if auth == nil {
		return nil
	}

	// generate access token
	ret := &AccessRequest{
		Type:            CLIENT_CREDENTIALS,
		Scope:           r.Form.Get("scope"),
		GenerateRefresh: true,
		Expiration:      s.Config.AccessExpiration,
	}

	// must have a valid client
	if ret.Client = getClient(auth, s.Storage, w); ret.Client == nil {
		return nil
	}

	// set redirect uri
	ret.RedirectUri = ret.Client.RedirectUri

	return ret
}

func (s *Server) FinishAccessRequest(w *Response, r *http.Request, ar *AccessRequest) {
	// don't process if is already an error
	if w.IsError {
		return
	}

	if ar.Authorized {
		// generate access token
		ret := &AccessData{
			Client:        ar.Client,
			AuthorizeData: ar.AuthorizeData,
			AccessData:    ar.AccessData,
			RedirectUri:   r.Form.Get("redirect_uri"),
			CreatedAt:     time.Now(),
			ExpiresIn:     ar.Expiration,
			UserData:      ar.UserData,
			Scope:         ar.Scope,
		}

		var err error

		// generate access token
		ret.AccessToken, ret.RefreshToken, err = s.AccessTokenGen.GenerateAccessToken(ret, ar.GenerateRefresh)
		if err != nil {
			w.SetError(E_SERVER_ERROR, "")
			w.InternalError = err
			return
		}

		// save access token
		if err = s.Storage.SaveAccess(ret); err != nil {
			w.SetError(E_SERVER_ERROR, "")
			w.InternalError = err
			return
		}

		// remove authorization token
		if ret.AuthorizeData != nil {
			s.Storage.RemoveAuthorize(ret.AuthorizeData.Code)
		}

		// remove previous access token
		if ret.AccessData != nil {
			if ret.AccessData.RefreshToken != "" {
				s.Storage.RemoveRefresh(ret.AccessData.RefreshToken)
			}
			s.Storage.RemoveAccess(ret.AccessData.AccessToken)
		}

		// output data
		w.Output["access_token"] = ret.AccessToken
		w.Output["token_type"] = s.Config.TokenType
		w.Output["expires_in"] = ret.ExpiresIn
		if ret.RefreshToken != "" {
			w.Output["refresh_token"] = ret.RefreshToken
		}
		if ar.Scope != "" {
			w.Output["scope"] = ar.Scope
		}
	} else {
		w.SetError(E_ACCESS_DENIED, "")
	}
}

// Helper Functions

// getClient looks up and authenticates the basic auth using the given
// storage. Sets an error on the response if auth fails or a server error occurs.
func getClient(auth *BasicAuth, storage Storage, w *Response) *Client {
	client, err := storage.GetClient(auth.Username)
	if err != nil {
		w.SetError(E_SERVER_ERROR, "")
		w.InternalError = err
		return nil
	}
	if client == nil {
		w.SetError(E_UNAUTHORIZED_CLIENT, "")
		return nil
	}
	if client.Secret != auth.Password {
		w.SetError(E_UNAUTHORIZED_CLIENT, "")
		return nil
	}
	if client.RedirectUri == "" {
		w.SetError(E_UNAUTHORIZED_CLIENT, "")
		return nil
	}
	return client
}
<commit_msg>UserData is copied from AuthorizeData to AccessRequest.<commit_after>package osin

import (
	"net/http"
	"time"
)

// AccessRequestType is the type for OAuth param `grant_type`
type AccessRequestType string

const (
	AUTHORIZATION_CODE AccessRequestType = "authorization_code"
	REFRESH_TOKEN                        = "refresh_token"
	PASSWORD                             = "password"
	CLIENT_CREDENTIALS                   = "client_credentials"
	IMPLICIT                             = "__implicit"
)

// AccessRequest is a request for access tokens
type AccessRequest struct {
	Type          AccessRequestType
	Code          string
	Client        *Client
	AuthorizeData *AuthorizeData
	AccessData    *AccessData
	RedirectUri   string
	Scope         string
	Username      string
	Password      string

	// Set if request is authorized
	Authorized bool

	// Token expiration in seconds. Change if different from default
	Expiration int32

	// Set if a refresh token should be generated
	GenerateRefresh bool

	// Data to be passed to storage. Not used by the library.
	UserData interface{}
}

// AccessData represents an access grant (tokens, expiration, client, etc)
type AccessData struct {
	// Client information
	Client *Client

	// Authorize data, for authorization code
	AuthorizeData *AuthorizeData

	// Previous access data, for refresh token
	AccessData *AccessData

	// Access token
	AccessToken string

	// Refresh Token. Can be blank
	RefreshToken string

	// Token expiration in seconds
	ExpiresIn int32

	// Requested scope
	Scope string

	// Redirect Uri from request
	RedirectUri string

	// Date created
	CreatedAt time.Time

	// Data to be passed to storage. Not used by the library.
	UserData interface{}
}

// IsExpired returns true if access expired
func (d *AccessData) IsExpired() bool {
	return d.CreatedAt.Add(time.Duration(d.ExpiresIn) * time.Second).Before(time.Now())
}

// ExpireAt returns the expiration date
func (d *AccessData) ExpireAt() time.Time {
	return d.CreatedAt.Add(time.Duration(d.ExpiresIn) * time.Second)
}

// AccessTokenGen generates access tokens
type AccessTokenGen interface {
	GenerateAccessToken(data *AccessData, generaterefresh bool) (accesstoken string, refreshtoken string, err error)
}

// HandleAccessRequest is the http.HandlerFunc for handling access token requests
func (s *Server) HandleAccessRequest(w *Response, r *http.Request) *AccessRequest {
	// Only allow GET or POST
	if r.Method == "GET" {
		if !s.Config.AllowGetAccessRequest {
			w.SetError(E_INVALID_REQUEST, "")
			return nil
		}
	} else if r.Method != "POST" {
		w.SetError(E_INVALID_REQUEST, "")
		return nil
	}

	err := r.ParseForm()
	if err != nil {
		w.SetError(E_INVALID_REQUEST, "")
		return nil
	}

	grantType := AccessRequestType(r.Form.Get("grant_type"))
	if s.Config.AllowedAccessTypes.Exists(grantType) {
		switch grantType {
		case AUTHORIZATION_CODE:
			return s.handleAuthorizationCodeRequest(w, r)
		case REFRESH_TOKEN:
			return s.handleRefreshTokenRequest(w, r)
		case PASSWORD:
			return s.handlePasswordRequest(w, r)
		case CLIENT_CREDENTIALS:
			return s.handleClientCredentialsRequest(w, r)
		}
	}

	w.SetError(E_UNSUPPORTED_GRANT_TYPE, "")
	return nil
}

func (s *Server) handleAuthorizationCodeRequest(w *Response, r *http.Request) *AccessRequest {
	// get client authentication
	auth := getClientAuth(w, r, s.Config.AllowClientSecretInParams)
	if auth == nil {
		return nil
	}

	// generate access token
	ret := &AccessRequest{
		Type:            AUTHORIZATION_CODE,
		Code:            r.Form.Get("code"),
		RedirectUri:     r.Form.Get("redirect_uri"),
		GenerateRefresh: true,
		Expiration:      s.Config.AccessExpiration,
	}

	// "code" is required
	if ret.Code == "" {
		w.SetError(E_INVALID_GRANT, "")
		return nil
	}

	// must have a valid client
	if ret.Client = getClient(auth, s.Storage, w); ret.Client == nil {
		return nil
	}

	// must be a valid authorization code
	var err error
	ret.AuthorizeData, err = s.Storage.LoadAuthorize(ret.Code)
	if err != nil {
		w.SetError(E_INVALID_GRANT, "")
		w.InternalError = err
		return nil
	}
	if ret.AuthorizeData.Client == nil {
		w.SetError(E_UNAUTHORIZED_CLIENT, "")
		return nil
	}
	if ret.AuthorizeData.Client.RedirectUri == "" {
		w.SetError(E_UNAUTHORIZED_CLIENT, "")
		return nil
	}
	if ret.AuthorizeData.IsExpired() {
		w.SetError(E_INVALID_GRANT, "")
		return nil
	}

	// code must be from the client
	if ret.AuthorizeData.Client.Id != ret.Client.Id {
		w.SetError(E_INVALID_GRANT, "")
		return nil
	}

	// check redirect uri
	if ret.RedirectUri == "" {
		ret.RedirectUri = ret.Client.RedirectUri
	}
	if err = ValidateUri(ret.Client.RedirectUri, ret.RedirectUri); err != nil {
		w.SetError(E_INVALID_REQUEST, "")
		w.InternalError = err
		return nil
	}
	if ret.AuthorizeData.RedirectUri != ret.RedirectUri {
		w.SetError(E_INVALID_REQUEST, "")
		return nil
	}

	// set rest of data
	ret.Scope = ret.AuthorizeData.Scope
	ret.UserData = ret.AuthorizeData.UserData

	return ret
}

func (s *Server) handleRefreshTokenRequest(w *Response, r *http.Request) *AccessRequest {
	// get client authentication
	auth := getClientAuth(w, r, s.Config.AllowClientSecretInParams)
	if auth == nil {
		return nil
	}

	// generate access token
	ret := &AccessRequest{
		Type:            REFRESH_TOKEN,
		Code:            r.Form.Get("refresh_token"),
		Scope:           r.Form.Get("scope"),
		GenerateRefresh: true,
		Expiration:      s.Config.AccessExpiration,
	}

	// "refresh_token" is required
	if ret.Code == "" {
		w.SetError(E_INVALID_GRANT, "")
		return nil
	}

	// must have a valid client
	if ret.Client = getClient(auth, s.Storage, w); ret.Client == nil {
		return nil
	}

	// must be a valid refresh code
	var err error
	ret.AccessData, err = s.Storage.LoadRefresh(ret.Code)
	if err != nil {
		w.SetError(E_INVALID_GRANT, "")
		w.InternalError = err
		return nil
	}
	if ret.AccessData.Client == nil {
		w.SetError(E_UNAUTHORIZED_CLIENT, "")
		return nil
	}
	if ret.AccessData.Client.RedirectUri == "" {
		w.SetError(E_UNAUTHORIZED_CLIENT, "")
		return nil
	}

	// client must be the same as the previous token
	if ret.AccessData.Client.Id != ret.Client.Id {
		w.SetError(E_INVALID_CLIENT, "")
		return nil

	}

	// set rest of data
	ret.RedirectUri = ret.AccessData.RedirectUri
	ret.UserData = ret.AccessData.UserData
	if ret.Scope == "" {
		ret.Scope = ret.AccessData.Scope
	}

	return ret
}

func (s *Server) handlePasswordRequest(w *Response, r *http.Request) *AccessRequest {
	// get client authentication
	auth := getClientAuth(w, r, s.Config.AllowClientSecretInParams)
	if auth == nil {
		return nil
	}

	// generate access token
	ret := &AccessRequest{
		Type:            PASSWORD,
		Username:        r.Form.Get("username"),
		Password:        r.Form.Get("password"),
		Scope:           r.Form.Get("scope"),
		GenerateRefresh: true,
		Expiration:      s.Config.AccessExpiration,
	}

	// "username" and "password" is required
	if ret.Username == "" || ret.Password == "" {
		w.SetError(E_INVALID_GRANT, "")
		return nil
	}

	// must have a valid client
	if ret.Client = getClient(auth, s.Storage, w); ret.Client == nil {
		return nil
	}

	// set redirect uri
	ret.RedirectUri = ret.Client.RedirectUri

	return ret
}

func (s *Server) handleClientCredentialsRequest(w *Response, r *http.Request) *AccessRequest {
	// get client authentication
	auth := getClientAuth(w, r, s.Config.AllowClientSecretInParams)
	if auth == nil {
		return nil
	}

	// generate access token
	ret := &AccessRequest{
		Type:            CLIENT_CREDENTIALS,
		Scope:           r.Form.Get("scope"),
		GenerateRefresh: true,
		Expiration:      s.Config.AccessExpiration,
	}

	// must have a valid client
	if ret.Client = getClient(auth, s.Storage, w); ret.Client == nil {
		return nil
	}

	// set redirect uri
	ret.RedirectUri = ret.Client.RedirectUri

	return ret
}

func (s *Server) FinishAccessRequest(w *Response, r *http.Request, ar *AccessRequest) {
	// don't process if is already an error
	if w.IsError {
		return
	}

	if ar.Authorized {
		// generate access token
		ret := &AccessData{
			Client:        ar.Client,
			AuthorizeData: ar.AuthorizeData,
			AccessData:    ar.AccessData,
			RedirectUri:   r.Form.Get("redirect_uri"),
			CreatedAt:     time.Now(),
			ExpiresIn:     ar.Expiration,
			UserData:      ar.UserData,
			Scope:         ar.Scope,
		}

		var err error

		// generate access token
		ret.AccessToken, ret.RefreshToken, err = s.AccessTokenGen.GenerateAccessToken(ret, ar.GenerateRefresh)
		if err != nil {
			w.SetError(E_SERVER_ERROR, "")
			w.InternalError = err
			return
		}

		// save access token
		if err = s.Storage.SaveAccess(ret); err != nil {
			w.SetError(E_SERVER_ERROR, "")
			w.InternalError = err
			return
		}

		// remove authorization token
		if ret.AuthorizeData != nil {
			s.Storage.RemoveAuthorize(ret.AuthorizeData.Code)
		}

		// remove previous access token
		if ret.AccessData != nil {
			if ret.AccessData.RefreshToken != "" {
				s.Storage.RemoveRefresh(ret.AccessData.RefreshToken)
			}
			s.Storage.RemoveAccess(ret.AccessData.AccessToken)
		}

		// output data
		w.Output["access_token"] = ret.AccessToken
		w.Output["token_type"] = s.Config.TokenType
		w.Output["expires_in"] = ret.ExpiresIn
		if ret.RefreshToken != "" {
			w.Output["refresh_token"] = ret.RefreshToken
		}
		if ar.Scope != "" {
			w.Output["scope"] = ar.Scope
		}
	} else {
		w.SetError(E_ACCESS_DENIED, "")
	}
}

// Helper Functions

// getClient looks up and authenticates the basic auth using the given
// storage. Sets an error on the response if auth fails or a server error occurs.
func getClient(auth *BasicAuth, storage Storage, w *Response) *Client {
	client, err := storage.GetClient(auth.Username)
	if err != nil {
		w.SetError(E_SERVER_ERROR, "")
		w.InternalError = err
		return nil
	}
	if client == nil {
		w.SetError(E_UNAUTHORIZED_CLIENT, "")
		return nil
	}
	if client.Secret != auth.Password {
		w.SetError(E_UNAUTHORIZED_CLIENT, "")
		return nil
	}
	if client.RedirectUri == "" {
		w.SetError(E_UNAUTHORIZED_CLIENT, "")
		return nil
	}
	return client
}
<|endoftext|>