<commit_before>package main

import (
	"fmt"
	"net"

	"github.com/fiorix/go-diameter/base"
	diam "github.com/fiorix/go-diameter/stack"
)

func main() {
	diam.HandleFunc(257, OnCER)
	diam.ListenAndServe(":3868", nil, nil)
}

// On CER reply with CEA.
// http://tools.ietf.org/html/rfc6733#section-5.3.2
func OnCER(w diam.ResponseWriter, r *diam.Request) {
	fmt.Println("Request from", r.RemoteAddr, "to", r.LocalAddr, ":")
	r.Msg.PrettyPrint()
	m := base.NewMessage(
		257,  // Command Code
		0x00, // Flags
		0x00, // Application Id
		r.Msg.Header.HopByHopId,
		r.Msg.Header.EndToEndId,
		r.Msg.Dict,
	)
	// Build message by attaching multiple AVPs to it.
	m.NewAVP(
		r.Dict.CodeFor("Result-Code"), 0x40, 0x0,
		&base.Unsigned32{Value: 2001},
	)
	m.NewAVP(
		r.Dict.CodeFor("Origin-Host"), 0x40, 0x0,
		&base.DiameterIdentity{base.OctetString{Value: "go"}},
	)
	m.NewAVP(
		r.Dict.CodeFor("Origin-Realm"), 0x40, 0x0,
		&base.DiameterIdentity{base.OctetString{Value: "server"}},
	)
	localIP, _, _ := net.SplitHostPort(r.LocalAddr)
	m.NewAVP(
		r.Dict.CodeFor("Host-IP-Address"), 0x40, 0x0,
		&base.Address{Family: []byte("01"), IP: net.ParseIP(localIP)},
	)
	m.NewAVP(
		r.Dict.CodeFor("Vendor-Id"), 0x40, 0x0,
		&base.Unsigned32{Value: 131313},
	)
	m.NewAVP(
		r.Dict.CodeFor("Product-Name"), 0x40, 0x0,
		&base.OctetString{Value: "go-diameter"},
	)
	// Reply with the same Origin-State-Id
	code := r.Dict.CodeFor("Origin-State-Id")
	OriginStateId, err := r.Msg.FindAVP(code)
	if err != nil {
		fmt.Println("Err:", err)
		return
	}
	m.NewAVP(code, 0x40, 0x0, OriginStateId.Body())

	// Write response
	fmt.Println("Response:")
	m.PrettyPrint()
	w.Write(m)
}
<commit_msg>Fixes<commit_after>package main

import (
	"fmt"
	"net"

	"github.com/fiorix/go-diameter/base"
	diam "github.com/fiorix/go-diameter/stack"
)

func main() {
	diam.HandleFunc(257, OnCER)
	diam.ListenAndServe(":3868", nil, nil)
}

// On CER reply with CEA.
// http://tools.ietf.org/html/rfc6733#section-5.3.2
func OnCER(w diam.ResponseWriter, r *diam.Request) {
	fmt.Println("Request from", r.RemoteAddr, "to", r.LocalAddr, ":")
	r.Msg.PrettyPrint()
	m := base.NewMessage(
		257,  // Command Code
		0x00, // Flags
		0x00, // Application Id
		r.Msg.Header.HopByHopId,
		r.Msg.Header.EndToEndId,
		r.Msg.Dict,
	)
	// Build message by attaching multiple AVPs to it.
	m.Add(base.NewAVP(r.Dict, r.Dict.CodeFor("Result-Code"), 0x40, 0x0,
		&base.Unsigned32{Value: 2001}))
	m.Add(base.NewAVP(r.Dict, r.Dict.CodeFor("Origin-Host"), 0x40, 0x0,
		&base.DiameterIdentity{base.OctetString{Value: "go"}}))
	m.Add(base.NewAVP(r.Dict, r.Dict.CodeFor("Origin-Realm"), 0x40, 0x0,
		&base.DiameterIdentity{base.OctetString{Value: "server"}}))
	localIP, _, _ := net.SplitHostPort(r.LocalAddr)
	m.Add(base.NewAVP(r.Dict, r.Dict.CodeFor("Host-IP-Address"), 0x40, 0x0,
		&base.Address{Family: []byte("01"), IP: net.ParseIP(localIP)}))
	m.Add(base.NewAVP(r.Dict, r.Dict.CodeFor("Vendor-Id"), 0x40, 0x0,
		&base.Unsigned32{Value: 131313}))
	m.Add(base.NewAVP(r.Dict, r.Dict.CodeFor("Product-Name"), 0x40, 0x0,
		&base.OctetString{Value: "go-diameter"}))
	// Reply with the same Origin-State-Id
	code := r.Dict.CodeFor("Origin-State-Id")
	OriginStateId, err := r.Msg.FindAVP(code)
	if err != nil {
		fmt.Println("Err:", err)
		return
	}
	m.Add(base.NewAVP(r.Dict, code, 0x40, 0x0, OriginStateId.Body()))

	// Write response
	fmt.Println("Response:")
	m.PrettyPrint()
	w.Write(m)
}
<|endoftext|>


<commit_before>package logrus_test

import (
	"bytes"
	"encoding/json"
	"sync"
	"testing"

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

	. "github.com/sirupsen/logrus"
	. "github.com/sirupsen/logrus/internal/testutils"
)

type TestHook struct {
	Fired bool
}

func (hook *TestHook) Fire(entry *Entry) error {
	hook.Fired = true
	return nil
}

func (hook *TestHook) Levels() []Level {
	return []Level{
		TraceLevel,
		DebugLevel,
		InfoLevel,
		WarnLevel,
		ErrorLevel,
		FatalLevel,
		PanicLevel,
	}
}

func TestHookFires(t *testing.T) {
	hook := new(TestHook)

	LogAndAssertJSON(t, func(log *Logger) {
		log.Hooks.Add(hook)
		assert.Equal(t, hook.Fired, false)

		log.Print("test")
	}, func(fields Fields) {
		assert.Equal(t, hook.Fired, true)
	})
}

type ModifyHook struct {
}

func (hook *ModifyHook) Fire(entry *Entry) error {
	entry.Data["wow"] = "whale"
	return nil
}

func (hook *ModifyHook) Levels() []Level {
	return []Level{
		TraceLevel,
		DebugLevel,
		InfoLevel,
		WarnLevel,
		ErrorLevel,
		FatalLevel,
		PanicLevel,
	}
}

func TestHookCanModifyEntry(t *testing.T) {
	hook := new(ModifyHook)

	LogAndAssertJSON(t, func(log *Logger) {
		log.Hooks.Add(hook)
		log.WithField("wow", "elephant").Print("test")
	}, func(fields Fields) {
		assert.Equal(t, fields["wow"], "whale")
	})
}

func TestCanFireMultipleHooks(t *testing.T) {
	hook1 := new(ModifyHook)
	hook2 := new(TestHook)

	LogAndAssertJSON(t, func(log *Logger) {
		log.Hooks.Add(hook1)
		log.Hooks.Add(hook2)

		log.WithField("wow", "elephant").Print("test")
	}, func(fields Fields) {
		assert.Equal(t, fields["wow"], "whale")
		assert.Equal(t, hook2.Fired, true)
	})
}

type SingleLevelModifyHook struct {
	ModifyHook
}

func (h *SingleLevelModifyHook) Levels() []Level {
	return []Level{InfoLevel}
}

func TestHookEntryIsPristine(t *testing.T) {
	l := New()
	b := &bytes.Buffer{}
	l.Formatter = &JSONFormatter{}
	l.Out = b
	l.AddHook(&SingleLevelModifyHook{})

	l.Error("error message")
	data := map[string]string{}
	err := json.Unmarshal(b.Bytes(), &data)
	require.NoError(t, err)
	_, ok := data["wow"]
	require.False(t, ok)
	b.Reset()

	l.Info("error message")
	data = map[string]string{}
	err = json.Unmarshal(b.Bytes(), &data)
	require.NoError(t, err)
	_, ok = data["wow"]
	require.True(t, ok)
	b.Reset()

	l.Error("error message")
	data = map[string]string{}
	err = json.Unmarshal(b.Bytes(), &data)
	require.NoError(t, err)
	_, ok = data["wow"]
	require.False(t, ok)
	b.Reset()
}

type ErrorHook struct {
	Fired bool
}

func (hook *ErrorHook) Fire(entry *Entry) error {
	hook.Fired = true
	return nil
}

func (hook *ErrorHook) Levels() []Level {
	return []Level{
		ErrorLevel,
	}
}

func TestErrorHookShouldntFireOnInfo(t *testing.T) {
	hook := new(ErrorHook)

	LogAndAssertJSON(t, func(log *Logger) {
		log.Hooks.Add(hook)
		log.Info("test")
	}, func(fields Fields) {
		assert.Equal(t, hook.Fired, false)
	})
}

func TestErrorHookShouldFireOnError(t *testing.T) {
	hook := new(ErrorHook)

	LogAndAssertJSON(t, func(log *Logger) {
		log.Hooks.Add(hook)
		log.Error("test")
	}, func(fields Fields) {
		assert.Equal(t, hook.Fired, true)
	})
}

func TestAddHookRace(t *testing.T) {
	var wg sync.WaitGroup
	wg.Add(2)
	hook := new(ErrorHook)
	LogAndAssertJSON(t, func(log *Logger) {
		go func() {
			defer wg.Done()
			log.AddHook(hook)
		}()
		go func() {
			defer wg.Done()
			log.Error("test")
		}()
		wg.Wait()
	}, func(fields Fields) {
		// the line may have been logged
		// before the hook was added, so we can't
		// actually assert on the hook
	})
}
<commit_msg>Add a unit test to ensure hook are called in their registration order<commit_after>package logrus_test

import (
	"bytes"
	"encoding/json"
	"sync"
	"testing"

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

	. "github.com/sirupsen/logrus"
	. "github.com/sirupsen/logrus/internal/testutils"
)

type TestHook struct {
	Fired bool
}

func (hook *TestHook) Fire(entry *Entry) error {
	hook.Fired = true
	return nil
}

func (hook *TestHook) Levels() []Level {
	return []Level{
		TraceLevel,
		DebugLevel,
		InfoLevel,
		WarnLevel,
		ErrorLevel,
		FatalLevel,
		PanicLevel,
	}
}

func TestHookFires(t *testing.T) {
	hook := new(TestHook)

	LogAndAssertJSON(t, func(log *Logger) {
		log.Hooks.Add(hook)
		assert.Equal(t, hook.Fired, false)

		log.Print("test")
	}, func(fields Fields) {
		assert.Equal(t, hook.Fired, true)
	})
}

type ModifyHook struct {
}

func (hook *ModifyHook) Fire(entry *Entry) error {
	entry.Data["wow"] = "whale"
	return nil
}

func (hook *ModifyHook) Levels() []Level {
	return []Level{
		TraceLevel,
		DebugLevel,
		InfoLevel,
		WarnLevel,
		ErrorLevel,
		FatalLevel,
		PanicLevel,
	}
}

func TestHookCanModifyEntry(t *testing.T) {
	hook := new(ModifyHook)

	LogAndAssertJSON(t, func(log *Logger) {
		log.Hooks.Add(hook)
		log.WithField("wow", "elephant").Print("test")
	}, func(fields Fields) {
		assert.Equal(t, fields["wow"], "whale")
	})
}

func TestCanFireMultipleHooks(t *testing.T) {
	hook1 := new(ModifyHook)
	hook2 := new(TestHook)

	LogAndAssertJSON(t, func(log *Logger) {
		log.Hooks.Add(hook1)
		log.Hooks.Add(hook2)

		log.WithField("wow", "elephant").Print("test")
	}, func(fields Fields) {
		assert.Equal(t, fields["wow"], "whale")
		assert.Equal(t, hook2.Fired, true)
	})
}

type SingleLevelModifyHook struct {
	ModifyHook
}

func (h *SingleLevelModifyHook) Levels() []Level {
	return []Level{InfoLevel}
}

func TestHookEntryIsPristine(t *testing.T) {
	l := New()
	b := &bytes.Buffer{}
	l.Formatter = &JSONFormatter{}
	l.Out = b
	l.AddHook(&SingleLevelModifyHook{})

	l.Error("error message")
	data := map[string]string{}
	err := json.Unmarshal(b.Bytes(), &data)
	require.NoError(t, err)
	_, ok := data["wow"]
	require.False(t, ok)
	b.Reset()

	l.Info("error message")
	data = map[string]string{}
	err = json.Unmarshal(b.Bytes(), &data)
	require.NoError(t, err)
	_, ok = data["wow"]
	require.True(t, ok)
	b.Reset()

	l.Error("error message")
	data = map[string]string{}
	err = json.Unmarshal(b.Bytes(), &data)
	require.NoError(t, err)
	_, ok = data["wow"]
	require.False(t, ok)
	b.Reset()
}

type ErrorHook struct {
	Fired bool
}

func (hook *ErrorHook) Fire(entry *Entry) error {
	hook.Fired = true
	return nil
}

func (hook *ErrorHook) Levels() []Level {
	return []Level{
		ErrorLevel,
	}
}

func TestErrorHookShouldntFireOnInfo(t *testing.T) {
	hook := new(ErrorHook)

	LogAndAssertJSON(t, func(log *Logger) {
		log.Hooks.Add(hook)
		log.Info("test")
	}, func(fields Fields) {
		assert.Equal(t, hook.Fired, false)
	})
}

func TestErrorHookShouldFireOnError(t *testing.T) {
	hook := new(ErrorHook)

	LogAndAssertJSON(t, func(log *Logger) {
		log.Hooks.Add(hook)
		log.Error("test")
	}, func(fields Fields) {
		assert.Equal(t, hook.Fired, true)
	})
}

func TestAddHookRace(t *testing.T) {
	var wg sync.WaitGroup
	wg.Add(2)
	hook := new(ErrorHook)
	LogAndAssertJSON(t, func(log *Logger) {
		go func() {
			defer wg.Done()
			log.AddHook(hook)
		}()
		go func() {
			defer wg.Done()
			log.Error("test")
		}()
		wg.Wait()
	}, func(fields Fields) {
		// the line may have been logged
		// before the hook was added, so we can't
		// actually assert on the hook
	})
}

type HookCallFunc struct {
	F func()
}

func (h *HookCallFunc) Levels() []Level {
	return AllLevels
}

func (h *HookCallFunc) Fire(e *Entry) error {
	h.F()
	return nil
}

func TestHookFireOrder(t *testing.T) {
	checkers := []string{}
	h := LevelHooks{}
	h.Add(&HookCallFunc{F: func() { checkers = append(checkers, "first hook") }})
	h.Add(&HookCallFunc{F: func() { checkers = append(checkers, "second hook") }})
	h.Add(&HookCallFunc{F: func() { checkers = append(checkers, "third hook") }})

	h.Fire(InfoLevel, &Entry{})
	require.Equal(t, []string{"first hook", "second hook", "third hook"}, checkers)
}
<|endoftext|>


<commit_before>package assured

import (
	// "net/http"

	kitlog "github.com/go-kit/kit/log"
)

// Client
type Client struct {
	Logger kitlog.Logger
}

// NewClient creates a new go-rest-assured client with the given parameters
func NewClient(l kitlog.Logger) *Client {
	return &Client{
		Logger: l,
	}
}

// Run starts the go-rest-assured service through the client
// func (c *Client) Run() error {

// }

// // Close is used to close the running service
// func (c *Client) Close() {

// }

// // URL returns the URL needed to use the stubbed assured endpoints
// func (c *Client) URL() (string, error) {

// }

// // Given stubs an assured Call
// func (c *Client) Given(call *Call) error {
// 	//Status.HttpOK if not set on call
// 	//strip leading `/`? use url library to strip query and stuff
// }

// // Verify returns all of the calls made against a stubbed method and path
// func (c *Client) Verify(method, path string) ([]*Call, error) {

// }

// // Clear assured calls for a Method and Path
// func (c *Client) Clear(method, path string) error {

// }

// // Clear assured calls for a Method and Path
// func (c *Client) Clear(method, path string) error {

// }
<commit_msg>client up<commit_after>package assured

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"

	kitlog "github.com/go-kit/kit/log"
	"github.com/phayes/freeport"
)

// Client
type Client struct {
	logger     kitlog.Logger
	ctx        context.Context
	errc       chan error
	port       *int
	httpClient *http.Client
}

// NewClient creates a new go-rest-assured client with the given parameters
func NewClient(l kitlog.Logger, port *int) *Client {
	return &Client{
		logger:     l,
		ctx:        context.Background(),
		errc:       make(chan error),
		port:       port,
		httpClient: &http.Client{},
	}
}

// Port gets the port that go-rest-assured will run on, or an open port if not set
func (c *Client) Port() int {
	if c.port != nil {
		return *c.port
	}
	port := freeport.GetPort()
	c.port = &port
	return port
}

// Run starts the go-rest-assured service through the client
func (c *Client) Run() {
	StartApplicationHTTPListener(c.Port(), c.logger, c.ctx, c.errc)
}

// Close is used to close the running service
func (c *Client) Close() {

}

// Given stubs an assured Call
func (c *Client) Given(call *Call) error {
	var req *http.Request
	var err error

	if call.Response == nil {
		req, err = http.NewRequest(call.Method, fmt.Sprintf("given/%s", call.Path), nil)
	} else {
		req, err = http.NewRequest(call.Method, fmt.Sprintf("given/%s", call.Path), bytes.NewReader(call.Response))
	}
	if err != nil {
		return err
	}
	if call.StatusCode != 0 {
		req.Header.Set("Assert-Status", string(call.StatusCode))
	}

	_, err = c.httpClient.Do(req)
	return err
}

// Verify returns all of the calls made against a stubbed method and path
func (c *Client) Verify(method, path string) ([]*Call, error) {
	req, err := http.NewRequest(method, fmt.Sprintf("verify/%s", path), nil)
	if err != nil {
		return nil, err
	}
	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("failure to verify calls")
	}
	defer resp.Body.Close()

	var calls []*Call
	if err = json.NewDecoder(resp.Body).Decode(calls); err != nil {
		return nil, err
	}
	return calls, nil
}

// Clear assured calls for a Method and Path
func (c *Client) Clear(method, path string) error {
	req, err := http.NewRequest(method, fmt.Sprintf("clear/%s", path), nil)
	if err != nil {
		return err
	}
	_, err = c.httpClient.Do(req)
	return err
}

// ClearAll clears all assured calls
func (c *Client) ClearAll() error {
	req, err := http.NewRequest(http.MethodDelete, "clear", nil)
	if err != nil {
		return err
	}
	_, err = c.httpClient.Do(req)
	return err
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"regexp"
	"sync"
	"time"

	pubsub "google.golang.org/api/pubsub/v1"

	log "github.com/Sirupsen/logrus"
)

type JobMessageStatus uint8

const (
	running JobMessageStatus = iota
	done
	acked
)

func (s JobMessageStatus) String() string {
	switch s {
	case running:
		return "running"
	case done:
		return "done"
	case acked:
		return "acked"
	default:
		return fmt.Sprintf("JobMessageStatus[Unknown:%v]", uint8(s))
	}
}

type (
	JobSustainerConfig struct {
		Disabled bool    `json:"disabled,omitempty"`
		Delay    float64 `json:"delay,omitempty"`
		Interval float64 `json:"interval,omitempty"`
	}

	JobMessage struct {
		sub    string
		raw    *pubsub.ReceivedMessage
		config *JobSustainerConfig
		puller Puller
		status JobMessageStatus
		mux    sync.Mutex
	}
)

func (m *JobMessage) Validate() error {
	if m.MessageId() == "" {
		return &InvalidJobError{msg: "no MessageId is given"}
	}
	return nil
}

func (m *JobMessage) MessageId() string {
	return m.raw.Message.MessageId
}

func (m *JobMessage) DownloadFiles() interface{} {
	str, ok := m.raw.Message.Attributes["download_files"]
	if !ok {
		return nil
	}
	return m.parseJson(str)
}

func (m *JobMessage) parseJson(str string) interface{} {
	matched, err := regexp.MatchString(`\A\[.*\]\z|\A\{.*\}\z|`, str)
	if err != nil {
		return str
	}
	if !matched {
		return str
	}
	var dest interface{}
	err = json.Unmarshal([]byte(str), &dest)
	if err != nil {
		return str
	}
	return dest
}

func (m *JobMessage) Ack() error {
	m.mux.Lock()
	defer m.mux.Unlock()

	logAttrs := log.Fields{"job_message_id": m.MessageId(), "ack_id": m.raw.AckId}
	log.WithFields(logAttrs).Debugln("Sending ACK")

	_, err := m.puller.Acknowledge(m.sub, m.raw.AckId)
	if err != nil {
		logAttrs["raw"] = fmt.Sprintf("%v", m.raw)
		logAttrs["error"] = err
		log.WithFields(logAttrs).Errorln("Failed to acknowledge")
		return err
	}

	logAttrs["status"] = m.status
	log.WithFields(logAttrs).Debugln("Updating status to acked")

	m.status = acked
	return nil
}

func (m *JobMessage) Nack() error {
	m.mux.Lock()
	defer m.mux.Unlock()

	logAttrs := log.Fields{"job_message_id": m.MessageId(), "ack_id": m.raw.AckId}
	log.WithFields(logAttrs).Debugln("Sending NACK")

	_, err := m.puller.ModifyAckDeadline(m.sub, []string{m.raw.AckId}, 0)
	if err != nil {
		logAttrs["raw"] = fmt.Sprintf("%v", m.raw)
		logAttrs["error"] = err
		log.WithFields(logAttrs).Errorln("Failed to send ModifyAckDeadline as a NACK")
		return err
	}

	logAttrs["status"] = m.status
	log.WithFields(logAttrs).Debugln("Updating status to done")

	m.status = done
	return nil
}

func (m *JobMessage) Done() {
	logAttrs := log.Fields{"job_message_id": m.MessageId(), "status": m.status}
	log.WithFields(logAttrs).Debugln("Done()")
	if m.status == running {
		m.status = done
	}
}

func (m *JobMessage) running() bool {
	return m.status == running
}

func (m *JobMessage) sendMADPeriodically(notification *ProgressNotification) error {
	if m.config.Disabled {
		return nil
	}
	log.Printf("sendMADPeriodically start\n")
	for {
		nextLimit := time.Now().Add(time.Duration(m.config.Interval) * time.Second)
		err := m.waitAndSendMAD(notification, nextLimit)
		if err != nil {
			log.WithFields(log.Fields{"error": err}).Errorln("Error in sendMADPeriodically")
			return err
		}
		if !m.running() {
			log.Debugln("sendMADPeriodically return")
			return nil
		}
	}
	// return nil
}

func (m *JobMessage) waitAndSendMAD(notification *ProgressNotification, nextLimit time.Time) error {
	log.Debugln("waitAndSendMAD starting")
	ticker := time.NewTicker(100 * time.Millisecond)
	for now := range ticker.C {
		if !m.running() {
			log.Debugln("waitAndSendMAD ticker stopping")
			ticker.Stop()
			log.Debugln("waitAndSendMAD ticker stopped")
			return nil
		}
		if now.After(nextLimit) {
			log.Debugln("waitAndSendMAD nextLimit passed")
			ticker.Stop()
			log.Debugln("waitAndSendMAD ticker stopped")
			break
		}
	}

	log.Debugln("waitAndSendMAD m.mux locking")

	m.mux.Lock()
	defer m.mux.Unlock()

	logAttrs := log.Fields{"status": m.status}
	log.WithFields(logAttrs).Debugln("waitAndSendMAD")

	// Don't send MAD after sending ACK
	if m.status == acked {
		log.WithFields(logAttrs).Infoln("waitAndSendMAD already acked")
		return nil
	}

	log.WithFields(logAttrs).Debugln("waitAndSendMAD sending ModifyAckDeadline")
	_, err := m.puller.ModifyAckDeadline(m.sub, []string{m.raw.AckId}, int64(m.config.Delay))
	if err != nil {
		logAttrs["error"] = err
		logAttrs["subscription"] = m.sub
		logAttrs["AckId"] = m.raw.AckId
		logAttrs["Delay"] = m.config.Delay
		log.WithFields(logAttrs).Errorln("waitAndSendMAD ModifyAckDeadline")
		msg := fmt.Sprintf("Failed modifyAckDeadline %v, %v, %v cause of %v\n", m.sub, m.raw.AckId, m.config.Delay, err)
		log.WithFields(logAttrs).Fatalf(msg)
		notification.notifyProgress(m.MessageId(), WORKING, false, log.ErrorLevel, msg)
	}
	return nil
}
<commit_msg>:+1: Implement support of Cloud Pub/Sub Notifications for Google Cloud Storage<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"regexp"
	"sync"
	"time"

	pubsub "google.golang.org/api/pubsub/v1"

	log "github.com/Sirupsen/logrus"
)

type JobMessageStatus uint8

const (
	running JobMessageStatus = iota
	done
	acked
)

func (s JobMessageStatus) String() string {
	switch s {
	case running:
		return "running"
	case done:
		return "done"
	case acked:
		return "acked"
	default:
		return fmt.Sprintf("JobMessageStatus[Unknown:%v]", uint8(s))
	}
}

type (
	JobSustainerConfig struct {
		Disabled bool    `json:"disabled,omitempty"`
		Delay    float64 `json:"delay,omitempty"`
		Interval float64 `json:"interval,omitempty"`
	}

	JobMessage struct {
		sub    string
		raw    *pubsub.ReceivedMessage
		config *JobSustainerConfig
		puller Puller
		status JobMessageStatus
		mux    sync.Mutex
	}
)

func (m *JobMessage) Validate() error {
	if m.MessageId() == "" {
		return &InvalidJobError{msg: "no MessageId is given"}
	}
	return nil
}

func (m *JobMessage) MessageId() string {
	return m.raw.Message.MessageId
}

func (m *JobMessage) DownloadFiles() interface{} {
	// See Cloud Pub/Sub Notifications for Google Cloud Storage
	// https://cloud.google.com/storage/docs/pubsub-notifications
	eventType, ok1 := m.raw.Message.Attributes["eventType"]
	bucketId, ok2 := m.raw.Message.Attributes["bucketId"]
	objectId, ok3 := m.raw.Message.Attributes["objectId"]
	if ok1 && ok2 && ok3 {
		if eventType != "OBJECT_FINALIZE" {
			return nil
		}
		url := "gs://" + bucketId + "/" + objectId
		return url
	}

	str, ok := m.raw.Message.Attributes["download_files"]
	if !ok {
		return nil
	}
	return m.parseJson(str)
}

func (m *JobMessage) parseJson(str string) interface{} {
	matched, err := regexp.MatchString(`\A\[.*\]\z|\A\{.*\}\z|`, str)
	if err != nil {
		return str
	}
	if !matched {
		return str
	}
	var dest interface{}
	err = json.Unmarshal([]byte(str), &dest)
	if err != nil {
		return str
	}
	return dest
}

func (m *JobMessage) Ack() error {
	m.mux.Lock()
	defer m.mux.Unlock()

	logAttrs := log.Fields{"job_message_id": m.MessageId(), "ack_id": m.raw.AckId}
	log.WithFields(logAttrs).Debugln("Sending ACK")

	_, err := m.puller.Acknowledge(m.sub, m.raw.AckId)
	if err != nil {
		logAttrs["raw"] = fmt.Sprintf("%v", m.raw)
		logAttrs["error"] = err
		log.WithFields(logAttrs).Errorln("Failed to acknowledge")
		return err
	}

	logAttrs["status"] = m.status
	log.WithFields(logAttrs).Debugln("Updating status to acked")

	m.status = acked
	return nil
}

func (m *JobMessage) Nack() error {
	m.mux.Lock()
	defer m.mux.Unlock()

	logAttrs := log.Fields{"job_message_id": m.MessageId(), "ack_id": m.raw.AckId}
	log.WithFields(logAttrs).Debugln("Sending NACK")

	_, err := m.puller.ModifyAckDeadline(m.sub, []string{m.raw.AckId}, 0)
	if err != nil {
		logAttrs["raw"] = fmt.Sprintf("%v", m.raw)
		logAttrs["error"] = err
		log.WithFields(logAttrs).Errorln("Failed to send ModifyAckDeadline as a NACK")
		return err
	}

	logAttrs["status"] = m.status
	log.WithFields(logAttrs).Debugln("Updating status to done")

	m.status = done
	return nil
}

func (m *JobMessage) Done() {
	logAttrs := log.Fields{"job_message_id": m.MessageId(), "status": m.status}
	log.WithFields(logAttrs).Debugln("Done()")
	if m.status == running {
		m.status = done
	}
}

func (m *JobMessage) running() bool {
	return m.status == running
}

func (m *JobMessage) sendMADPeriodically(notification *ProgressNotification) error {
	if m.config.Disabled {
		return nil
	}
	log.Printf("sendMADPeriodically start\n")
	for {
		nextLimit := time.Now().Add(time.Duration(m.config.Interval) * time.Second)
		err := m.waitAndSendMAD(notification, nextLimit)
		if err != nil {
			log.WithFields(log.Fields{"error": err}).Errorln("Error in sendMADPeriodically")
			return err
		}
		if !m.running() {
			log.Debugln("sendMADPeriodically return")
			return nil
		}
	}
	// return nil
}

func (m *JobMessage) waitAndSendMAD(notification *ProgressNotification, nextLimit time.Time) error {
	log.Debugln("waitAndSendMAD starting")
	ticker := time.NewTicker(100 * time.Millisecond)
	for now := range ticker.C {
		if !m.running() {
			log.Debugln("waitAndSendMAD ticker stopping")
			ticker.Stop()
			log.Debugln("waitAndSendMAD ticker stopped")
			return nil
		}
		if now.After(nextLimit) {
			log.Debugln("waitAndSendMAD nextLimit passed")
			ticker.Stop()
			log.Debugln("waitAndSendMAD ticker stopped")
			break
		}
	}

	log.Debugln("waitAndSendMAD m.mux locking")

	m.mux.Lock()
	defer m.mux.Unlock()

	logAttrs := log.Fields{"status": m.status}
	log.WithFields(logAttrs).Debugln("waitAndSendMAD")

	// Don't send MAD after sending ACK
	if m.status == acked {
		log.WithFields(logAttrs).Infoln("waitAndSendMAD already acked")
		return nil
	}

	log.WithFields(logAttrs).Debugln("waitAndSendMAD sending ModifyAckDeadline")
	_, err := m.puller.ModifyAckDeadline(m.sub, []string{m.raw.AckId}, int64(m.config.Delay))
	if err != nil {
		logAttrs["error"] = err
		logAttrs["subscription"] = m.sub
		logAttrs["AckId"] = m.raw.AckId
		logAttrs["Delay"] = m.config.Delay
		log.WithFields(logAttrs).Errorln("waitAndSendMAD ModifyAckDeadline")
		msg := fmt.Sprintf("Failed modifyAckDeadline %v, %v, %v cause of %v\n", m.sub, m.raw.AckId, m.config.Delay, err)
		log.WithFields(logAttrs).Fatalf(msg)
		notification.notifyProgress(m.MessageId(), WORKING, false, log.ErrorLevel, msg)
	}
	return nil
}
<|endoftext|>


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

// Package auth contains functions for minting custom authentication tokens, and verifying Firebase ID tokens.
package auth

import (
	"context"
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"strings"

	"firebase.google.com/go/internal"
	"google.golang.org/api/identitytoolkit/v3"
	"google.golang.org/api/transport"
)

const firebaseAudience = "https://identitytoolkit.googleapis.com/google.identity.identitytoolkit.v1.IdentityToolkit"
const googleCertURL = "https://www.googleapis.com/robot/v1/metadata/x509/securetoken@system.gserviceaccount.com"
const issuerPrefix = "https://securetoken.google.com/"
const tokenExpSeconds = 3600

var reservedClaims = []string{
	"acr", "amr", "at_hash", "aud", "auth_time", "azp", "cnf", "c_hash",
	"exp", "firebase", "iat", "iss", "jti", "nbf", "nonce", "sub",
}

var clk clock = &systemClock{}

// Token represents a decoded Firebase ID token.
//
// Token provides typed accessors to the common JWT fields such as Audience (aud) and Expiry (exp).
// Additionally it provides a UID field, which indicates the user ID of the account to which this token
// belongs. Any additional JWT claims can be accessed via the Claims map of Token.
type Token struct {
	Issuer   string                 `json:"iss"`
	Audience string                 `json:"aud"`
	Expires  int64                  `json:"exp"`
	IssuedAt int64                  `json:"iat"`
	Subject  string                 `json:"sub,omitempty"`
	UID      string                 `json:"uid,omitempty"`
	Claims   map[string]interface{} `json:"-"`
}

// Client is the interface for the Firebase auth service.
//
// Client facilitates generating custom JWT tokens for Firebase clients, and verifying ID tokens issued
// by Firebase backend services.
type Client struct {
	hc        *internal.HTTPClient
	is        *identitytoolkit.Service
	ks        keySource
	projectID string
	snr       signer
	version   string
}

type signer interface {
	Email() (string, error)
	Sign(b []byte) ([]byte, error)
}

// NewClient creates a new instance of the Firebase Auth Client.
//
// This function can only be invoked from within the SDK. Client applications should access the
// Auth service through firebase.App.
func NewClient(ctx context.Context, c *internal.AuthConfig) (*Client, error) {
	var (
		err   error
		email string
		pk    *rsa.PrivateKey
	)
	if c.Creds != nil && len(c.Creds.JSON) > 0 {
		var svcAcct struct {
			ClientEmail string `json:"client_email"`
			PrivateKey  string `json:"private_key"`
		}
		if err := json.Unmarshal(c.Creds.JSON, &svcAcct); err != nil {
			return nil, err
		}
		if svcAcct.PrivateKey != "" {
			pk, err = parseKey(svcAcct.PrivateKey)
			if err != nil {
				return nil, err
			}
		}
		email = svcAcct.ClientEmail
	}

	var snr signer
	if email != "" && pk != nil {
		snr = serviceAcctSigner{email: email, pk: pk}
	} else {
		snr, err = newSigner(ctx)
		if err != nil {
			return nil, err
		}
	}

	hc, _, err := transport.NewHTTPClient(ctx, c.Opts...)
	if err != nil {
		return nil, err
	}

	is, err := identitytoolkit.New(hc)
	if err != nil {
		return nil, err
	}

	return &Client{
		hc:        &internal.HTTPClient{Client: hc},
		is:        is,
		ks:        newHTTPKeySource(googleCertURL, hc),
		projectID: c.ProjectID,
		snr:       snr,
		version:   "Go/Admin/" + c.Version,
	}, nil
}

// CustomToken creates a signed custom authentication token with the specified user ID. The resulting
// JWT can be used in a Firebase client SDK to trigger an authentication flow. See
// https://firebase.google.com/docs/auth/admin/create-custom-tokens#sign_in_using_custom_tokens_on_clients
// for more details on how to use custom tokens for client authentication.
func (c *Client) CustomToken(uid string) (string, error) {
	return c.CustomTokenWithClaims(uid, nil)
}

// CustomTokenWithClaims is similar to CustomToken, but in addition to the user ID, it also encodes
// all the key-value pairs in the provided map as claims in the resulting JWT.
func (c *Client) CustomTokenWithClaims(uid string, devClaims map[string]interface{}) (string, error) {
	iss, err := c.snr.Email()
	if err != nil {
		return "", err
	}

	if len(uid) == 0 || len(uid) > 128 {
		return "", errors.New("uid must be non-empty, and not longer than 128 characters")
	}

	var disallowed []string
	for _, k := range reservedClaims {
		if _, contains := devClaims[k]; contains {
			disallowed = append(disallowed, k)
		}
	}
	if len(disallowed) == 1 {
		return "", fmt.Errorf("developer claim %q is reserved and cannot be specified", disallowed[0])
	} else if len(disallowed) > 1 {
		return "", fmt.Errorf("developer claims %q are reserved and cannot be specified", strings.Join(disallowed, ", "))
	}

	now := clk.Now().Unix()
	payload := &customToken{
		Iss:    iss,
		Sub:    iss,
		Aud:    firebaseAudience,
		UID:    uid,
		Iat:    now,
		Exp:    now + tokenExpSeconds,
		Claims: devClaims,
	}
	return encodeToken(c.snr, defaultHeader(), payload)
}

// RevokeRefreshTokens revokes all refresh tokens issued to a user.
//
// RevokeRefreshTokens updates the user's TokensValidAfterMillis to the current UTC second.
// It is important that the server on which this is called has its clock set correctly and synchronized.
//
// While this revokes all sessions for a specified user and disables any new ID tokens for existing sessions
// from getting minted, existing ID tokens may remain active until their natural expiration (one hour).
// To verify that ID tokens are revoked, use `verifyIdTokenAndCheckRevoked(ctx, idToken)`.
func (c *Client) RevokeRefreshTokens(ctx context.Context, uid string) error {
	return c.updateUser(ctx, uid, (&UserToUpdate{}).revokeRefreshTokens())
}

// VerifyIDToken verifies the signature	and payload of the provided ID token.
//
// VerifyIDToken accepts a signed JWT token string, and verifies that it is current, issued for the
// correct Firebase project, and signed by the Google Firebase services in the cloud. It returns
// a Token containing the decoded claims in the input JWT. See
// https://firebase.google.com/docs/auth/admin/verify-id-tokens#retrieve_id_tokens_on_clients for
// more details on how to obtain an ID token in a client app.
// This does not check whether or not the token has been revoked. See `VerifyIDTokenAndCheckRevoked` below.
func (c *Client) VerifyIDToken(idToken string) (*Token, error) {
	if c.projectID == "" {
		return nil, errors.New("project id not available")
	}
	if idToken == "" {
		return nil, fmt.Errorf("ID token must be a non-empty string")
	}

	h := &jwtHeader{}
	p := &Token{}
	if err := decodeToken(idToken, c.ks, h, p); err != nil {
		return nil, err
	}

	projectIDMsg := "Make sure the ID token comes from the same Firebase project as the credential used to" +
		" authenticate this SDK."
	verifyTokenMsg := "See https://firebase.google.com/docs/auth/admin/verify-id-tokens for details on how to " +
		"retrieve a valid ID token."
	issuer := issuerPrefix + c.projectID

	var err error
	if h.KeyID == "" {
		if p.Audience == firebaseAudience {
			err = fmt.Errorf("VerifyIDToken() expects an ID token, but was given a custom token")
		} else {
			err = fmt.Errorf("ID token has no 'kid' header")
		}
	} else if h.Algorithm != "RS256" {
		err = fmt.Errorf("ID token has invalid incorrect algorithm. Expected 'RS256' but got %q. %s",
			h.Algorithm, verifyTokenMsg)
	} else if p.Audience != c.projectID {
		err = fmt.Errorf("ID token has invalid 'aud' (audience) claim. Expected %q but got %q. %s %s",
			c.projectID, p.Audience, projectIDMsg, verifyTokenMsg)
	} else if p.Issuer != issuer {
		err = fmt.Errorf("ID token has invalid 'iss' (issuer) claim. Expected %q but got %q. %s %s",
			issuer, p.Issuer, projectIDMsg, verifyTokenMsg)
	} else if p.IssuedAt > clk.Now().Unix() {
		err = fmt.Errorf("ID token issued at future timestamp: %d", p.IssuedAt)
	} else if p.Expires < clk.Now().Unix() {
		err = fmt.Errorf("ID token has expired. Expired at: %d", p.Expires)
	} else if p.Subject == "" {
		err = fmt.Errorf("ID token has empty 'sub' (subject) claim. %s", verifyTokenMsg)
	} else if len(p.Subject) > 128 {
		err = fmt.Errorf("ID token has a 'sub' (subject) claim longer than 128 characters. %s", verifyTokenMsg)
	}

	if err != nil {
		return nil, err
	}
	p.UID = p.Subject
	return p, nil
}

// VerifyIDTokenAndCheckRevoked verifies the provided ID token and checks it has not been revoked.
//
// VerifyIDTokenAndCheckRevoked verifies the signature and payload of the provided ID token and
// checks that it wasn't revoked. Uses VerifyIDToken() internally to verify the ID token JWT.
func (c *Client) VerifyIDTokenAndCheckRevoked(ctx context.Context, idToken string) (*Token, error) {
	p, err := c.VerifyIDToken(idToken)
	if err != nil {
		return nil, err
	}

	user, err := c.GetUser(ctx, p.UID)
	if err != nil {
		return nil, err
	}

	if p.IssuedAt*1000 < user.TokensValidAfterMillis {
		return nil, fmt.Errorf("ID token has been revoked")
	}
	return p, nil
}

func parseKey(key string) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode([]byte(key))
	if block == nil {
		return nil, fmt.Errorf("no private key data found in: %v", key)
	}
	k := block.Bytes
	parsedKey, err := x509.ParsePKCS8PrivateKey(k)
	if err != nil {
		parsedKey, err = x509.ParsePKCS1PrivateKey(k)
		if err != nil {
			return nil, fmt.Errorf("private key should be a PEM or plain PKSC1 or PKCS8; parse error: %v", err)
		}
	}
	parsed, ok := parsedKey.(*rsa.PrivateKey)
	if !ok {
		return nil, errors.New("private key is not an RSA key")
	}
	return parsed, nil
}
<commit_msg>Removing unused member from auth.Client (#118)<commit_after>// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package auth contains functions for minting custom authentication tokens, and verifying Firebase ID tokens.
package auth

import (
	"context"
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"strings"

	"firebase.google.com/go/internal"
	"google.golang.org/api/identitytoolkit/v3"
	"google.golang.org/api/transport"
)

const firebaseAudience = "https://identitytoolkit.googleapis.com/google.identity.identitytoolkit.v1.IdentityToolkit"
const googleCertURL = "https://www.googleapis.com/robot/v1/metadata/x509/securetoken@system.gserviceaccount.com"
const issuerPrefix = "https://securetoken.google.com/"
const tokenExpSeconds = 3600

var reservedClaims = []string{
	"acr", "amr", "at_hash", "aud", "auth_time", "azp", "cnf", "c_hash",
	"exp", "firebase", "iat", "iss", "jti", "nbf", "nonce", "sub",
}

var clk clock = &systemClock{}

// Token represents a decoded Firebase ID token.
//
// Token provides typed accessors to the common JWT fields such as Audience (aud) and Expiry (exp).
// Additionally it provides a UID field, which indicates the user ID of the account to which this token
// belongs. Any additional JWT claims can be accessed via the Claims map of Token.
type Token struct {
	Issuer   string                 `json:"iss"`
	Audience string                 `json:"aud"`
	Expires  int64                  `json:"exp"`
	IssuedAt int64                  `json:"iat"`
	Subject  string                 `json:"sub,omitempty"`
	UID      string                 `json:"uid,omitempty"`
	Claims   map[string]interface{} `json:"-"`
}

// Client is the interface for the Firebase auth service.
//
// Client facilitates generating custom JWT tokens for Firebase clients, and verifying ID tokens issued
// by Firebase backend services.
type Client struct {
	is        *identitytoolkit.Service
	ks        keySource
	projectID string
	snr       signer
	version   string
}

type signer interface {
	Email() (string, error)
	Sign(b []byte) ([]byte, error)
}

// NewClient creates a new instance of the Firebase Auth Client.
//
// This function can only be invoked from within the SDK. Client applications should access the
// Auth service through firebase.App.
func NewClient(ctx context.Context, c *internal.AuthConfig) (*Client, error) {
	var (
		err   error
		email string
		pk    *rsa.PrivateKey
	)
	if c.Creds != nil && len(c.Creds.JSON) > 0 {
		var svcAcct struct {
			ClientEmail string `json:"client_email"`
			PrivateKey  string `json:"private_key"`
		}
		if err := json.Unmarshal(c.Creds.JSON, &svcAcct); err != nil {
			return nil, err
		}
		if svcAcct.PrivateKey != "" {
			pk, err = parseKey(svcAcct.PrivateKey)
			if err != nil {
				return nil, err
			}
		}
		email = svcAcct.ClientEmail
	}

	var snr signer
	if email != "" && pk != nil {
		snr = serviceAcctSigner{email: email, pk: pk}
	} else {
		snr, err = newSigner(ctx)
		if err != nil {
			return nil, err
		}
	}

	hc, _, err := transport.NewHTTPClient(ctx, c.Opts...)
	if err != nil {
		return nil, err
	}

	is, err := identitytoolkit.New(hc)
	if err != nil {
		return nil, err
	}

	return &Client{
		is:        is,
		ks:        newHTTPKeySource(googleCertURL, hc),
		projectID: c.ProjectID,
		snr:       snr,
		version:   "Go/Admin/" + c.Version,
	}, nil
}

// CustomToken creates a signed custom authentication token with the specified user ID. The resulting
// JWT can be used in a Firebase client SDK to trigger an authentication flow. See
// https://firebase.google.com/docs/auth/admin/create-custom-tokens#sign_in_using_custom_tokens_on_clients
// for more details on how to use custom tokens for client authentication.
func (c *Client) CustomToken(uid string) (string, error) {
	return c.CustomTokenWithClaims(uid, nil)
}

// CustomTokenWithClaims is similar to CustomToken, but in addition to the user ID, it also encodes
// all the key-value pairs in the provided map as claims in the resulting JWT.
func (c *Client) CustomTokenWithClaims(uid string, devClaims map[string]interface{}) (string, error) {
	iss, err := c.snr.Email()
	if err != nil {
		return "", err
	}

	if len(uid) == 0 || len(uid) > 128 {
		return "", errors.New("uid must be non-empty, and not longer than 128 characters")
	}

	var disallowed []string
	for _, k := range reservedClaims {
		if _, contains := devClaims[k]; contains {
			disallowed = append(disallowed, k)
		}
	}
	if len(disallowed) == 1 {
		return "", fmt.Errorf("developer claim %q is reserved and cannot be specified", disallowed[0])
	} else if len(disallowed) > 1 {
		return "", fmt.Errorf("developer claims %q are reserved and cannot be specified", strings.Join(disallowed, ", "))
	}

	now := clk.Now().Unix()
	payload := &customToken{
		Iss:    iss,
		Sub:    iss,
		Aud:    firebaseAudience,
		UID:    uid,
		Iat:    now,
		Exp:    now + tokenExpSeconds,
		Claims: devClaims,
	}
	return encodeToken(c.snr, defaultHeader(), payload)
}

// RevokeRefreshTokens revokes all refresh tokens issued to a user.
//
// RevokeRefreshTokens updates the user's TokensValidAfterMillis to the current UTC second.
// It is important that the server on which this is called has its clock set correctly and synchronized.
//
// While this revokes all sessions for a specified user and disables any new ID tokens for existing sessions
// from getting minted, existing ID tokens may remain active until their natural expiration (one hour).
// To verify that ID tokens are revoked, use `verifyIdTokenAndCheckRevoked(ctx, idToken)`.
func (c *Client) RevokeRefreshTokens(ctx context.Context, uid string) error {
	return c.updateUser(ctx, uid, (&UserToUpdate{}).revokeRefreshTokens())
}

// VerifyIDToken verifies the signature	and payload of the provided ID token.
//
// VerifyIDToken accepts a signed JWT token string, and verifies that it is current, issued for the
// correct Firebase project, and signed by the Google Firebase services in the cloud. It returns
// a Token containing the decoded claims in the input JWT. See
// https://firebase.google.com/docs/auth/admin/verify-id-tokens#retrieve_id_tokens_on_clients for
// more details on how to obtain an ID token in a client app.
// This does not check whether or not the token has been revoked. See `VerifyIDTokenAndCheckRevoked` below.
func (c *Client) VerifyIDToken(idToken string) (*Token, error) {
	if c.projectID == "" {
		return nil, errors.New("project id not available")
	}
	if idToken == "" {
		return nil, fmt.Errorf("ID token must be a non-empty string")
	}

	h := &jwtHeader{}
	p := &Token{}
	if err := decodeToken(idToken, c.ks, h, p); err != nil {
		return nil, err
	}

	projectIDMsg := "Make sure the ID token comes from the same Firebase project as the credential used to" +
		" authenticate this SDK."
	verifyTokenMsg := "See https://firebase.google.com/docs/auth/admin/verify-id-tokens for details on how to " +
		"retrieve a valid ID token."
	issuer := issuerPrefix + c.projectID

	var err error
	if h.KeyID == "" {
		if p.Audience == firebaseAudience {
			err = fmt.Errorf("VerifyIDToken() expects an ID token, but was given a custom token")
		} else {
			err = fmt.Errorf("ID token has no 'kid' header")
		}
	} else if h.Algorithm != "RS256" {
		err = fmt.Errorf("ID token has invalid incorrect algorithm. Expected 'RS256' but got %q. %s",
			h.Algorithm, verifyTokenMsg)
	} else if p.Audience != c.projectID {
		err = fmt.Errorf("ID token has invalid 'aud' (audience) claim. Expected %q but got %q. %s %s",
			c.projectID, p.Audience, projectIDMsg, verifyTokenMsg)
	} else if p.Issuer != issuer {
		err = fmt.Errorf("ID token has invalid 'iss' (issuer) claim. Expected %q but got %q. %s %s",
			issuer, p.Issuer, projectIDMsg, verifyTokenMsg)
	} else if p.IssuedAt > clk.Now().Unix() {
		err = fmt.Errorf("ID token issued at future timestamp: %d", p.IssuedAt)
	} else if p.Expires < clk.Now().Unix() {
		err = fmt.Errorf("ID token has expired. Expired at: %d", p.Expires)
	} else if p.Subject == "" {
		err = fmt.Errorf("ID token has empty 'sub' (subject) claim. %s", verifyTokenMsg)
	} else if len(p.Subject) > 128 {
		err = fmt.Errorf("ID token has a 'sub' (subject) claim longer than 128 characters. %s", verifyTokenMsg)
	}

	if err != nil {
		return nil, err
	}
	p.UID = p.Subject
	return p, nil
}

// VerifyIDTokenAndCheckRevoked verifies the provided ID token and checks it has not been revoked.
//
// VerifyIDTokenAndCheckRevoked verifies the signature and payload of the provided ID token and
// checks that it wasn't revoked. Uses VerifyIDToken() internally to verify the ID token JWT.
func (c *Client) VerifyIDTokenAndCheckRevoked(ctx context.Context, idToken string) (*Token, error) {
	p, err := c.VerifyIDToken(idToken)
	if err != nil {
		return nil, err
	}

	user, err := c.GetUser(ctx, p.UID)
	if err != nil {
		return nil, err
	}

	if p.IssuedAt*1000 < user.TokensValidAfterMillis {
		return nil, fmt.Errorf("ID token has been revoked")
	}
	return p, nil
}

func parseKey(key string) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode([]byte(key))
	if block == nil {
		return nil, fmt.Errorf("no private key data found in: %v", key)
	}
	k := block.Bytes
	parsedKey, err := x509.ParsePKCS8PrivateKey(k)
	if err != nil {
		parsedKey, err = x509.ParsePKCS1PrivateKey(k)
		if err != nil {
			return nil, fmt.Errorf("private key should be a PEM or plain PKSC1 or PKCS8; parse error: %v", err)
		}
	}
	parsed, ok := parsedKey.(*rsa.PrivateKey)
	if !ok {
		return nil, errors.New("private key is not an RSA key")
	}
	return parsed, nil
}
<|endoftext|>


<commit_before>package auth

import (
	"errors"
	"fmt"
	"net/http"
	"strings"
	"time"

	"net/url"

	"github.com/chrisolsen/fbgraphapi"
	"golang.org/x/net/context"
	"google.golang.org/appengine"
	"google.golang.org/appengine/datastore"
)

// Errors
var (
	ErrNoCookie    = errors.New("no cookie found")
	ErrNoAuthToken = errors.New("no header auth token found")
)

// GetToken returns the *Token value for the raw token value contained within the auth cookie or auth header
func GetToken(c context.Context, r *http.Request) (*Token, error) {
	uuid, err := getUUIDFromCookie(r)
	if err != nil && err != http.ErrNoCookie {
		return nil, err
	}
	if err == http.ErrNoCookie {
		uuid = getUUIDFromHeader(r)
	}
	if err != nil {
		return nil, err
	}
	if len(uuid) == 0 {
		return nil, ErrNoAuthToken
	}

	tstore := NewTokenStore()
	return tstore.Get(c, uuid)
}

// Signup creates a user account and links up the credentials. Based on the request type an auth cookie
// or header token will be set with an auth token.
func Signup(c context.Context, w http.ResponseWriter, r *http.Request, creds *Credentials, account *Account) error {
	astore := NewAccountStore()
	tstore := NewTokenStore()

	accountKey, err := astore.Create(c, creds, account)
	if err != nil {
		return errors.New("failed to create account")
	}
	token, err := tstore.Create(c, accountKey)

	if accepts(r, "json") {
		setHeaderToken(w, token.UUID)
	} else {
		setAuthCookieToken(w, token.UUID)
	}

	return nil
}

// Signout deletes the token in the response to the client as well as deletes the token
// in the database to ensure it is no longer usable
func Signout(c context.Context, w http.ResponseWriter, r *http.Request) error {
	var err error
	var uuid string

	if accepts(r, "json") {
		uuid = getUUIDFromHeader(r)
		if err != nil {
			return fmt.Errorf("failed to get token from header: %v", err)
		}
		clearHeader(w)
	} else {
		uuid, err = getUUIDFromCookie(r)
		if err != nil {
			return fmt.Errorf("failed to get token from cookie: %v", err)
		}
		clearCookie(w)
	}

	store := NewTokenStore()
	token, err := store.Get(c, uuid)
	if err != nil {
		return err
	}
	err = datastore.Delete(c, token.Key)

	return err
}

// Authorize .
func Authorize(c context.Context, w http.ResponseWriter, r *http.Request, creds *Credentials) (*Token, error) {
	var token *Token
	var err error
	if len(creds.ProviderName) > 0 {
		token, err = authorizeOut(c, creds, appEngineURLGetter{Ctx: c})
	} else {
		token, err = authorizeIn(c, creds)
	}
	if err != nil {
		return nil, err
	}

	if accepts(r, "json") {
		setHeaderToken(w, token.UUID)
	} else {
		setAuthCookieToken(w, token.UUID)
	}

	return token, nil
}

func authorizeIn(c context.Context, creds *Credentials) (*Token, error) {
	accountStore := NewAccountStore()
	tokenStore := NewTokenStore()

	accountKey, err := accountStore.GetAccountKeyByCredentials(c, creds)
	if err != nil {
		return nil, fmt.Errorf("getting account key by credentials: %v", err)
	}

	token, err := tokenStore.Create(c, accountKey)
	if err != nil {
		return nil, err
	}

	return token, nil
}

func authorizeOut(c context.Context, creds *Credentials, urlGetter urlGetter) (*Token, error) {
	var err error
	tokenStore := NewTokenStore()
	accountStore := NewAccountStore()

	switch creds.ProviderName {
	case "facebook":
		err = fbgraphapi.Authenticate(creds.ProviderToken, creds.ProviderID, urlGetter)
	default:
		return nil, errors.New("unknown auth provider")
	}
	if err != nil {
		return nil, fmt.Errorf("authenticate: %v", err)
	}

	accountKey, err := accountStore.GetAccountKeyByCredentials(c, creds)
	if err != nil {
		return nil, fmt.Errorf("getting account key by credentials: %v", err)
	}

	token, err := tokenStore.Create(c, accountKey)
	if err != nil {
		return nil, err
	}

	return token, nil
}

func clearCookie(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     cookieName,
		Path:     "/",
		Expires:  time.Time{},
		HttpOnly: true,
		Value:    "",
		Secure:   !appengine.IsDevAppServer(),
	})
}

func setAuthCookieToken(w http.ResponseWriter, uuid string) {
	http.SetCookie(w, &http.Cookie{
		Name:     cookieName,
		Path:     "/",
		Expires:  time.Now().Add(time.Hour * 24 * 14),
		HttpOnly: true,
		Secure:   !appengine.IsDevAppServer(),
		Value:    uuid,
	})
}

func setHeaderToken(w http.ResponseWriter, uuid string) {
	w.Header().Set("Authorization", uuid)
}

func getUUIDFromHeader(r *http.Request) string {
	return r.Header.Get("Authorization")[len("token="):]
}

func getUUIDFromCookie(r *http.Request) (string, error) {
	cookie, err := r.Cookie(cookieName)
	if err != nil || len(cookie.Value) == 0 {
		return "", ErrNoCookie
	}

	return cookie.Value, nil
}

func clearHeader(w http.ResponseWriter) {
	w.Header().Del("Authorization")
}

func accepts(r *http.Request, t string) bool {
	return strings.Index(r.Header.Get("Accept"), t) > 0
}

// VerifyReferrer middlware validates the referer header matches the request url's host
func VerifyReferrer(c context.Context, w http.ResponseWriter, r *http.Request) context.Context {
	if r.Method != http.MethodPost {
		return c
	}

	err := func() error {
		referrer, err := url.Parse(r.Header.Get("Referer"))
		if err != nil {
			return err
		}
		if referrer.Host != r.Host {
			return fmt.Errorf("BLOCKED: VerifyReferrer: %v - %v", referrer, r.Host)
		}
		return nil
	}()

	if err != nil {
		cc, cancel := context.WithCancel(c)
		cancel()
		return cc
	}
	return c
}
<commit_msg>rename some methods<commit_after>package auth

import (
	"errors"
	"fmt"
	"net/http"
	"strings"
	"time"

	"net/url"

	"github.com/chrisolsen/fbgraphapi"
	"golang.org/x/net/context"
	"google.golang.org/appengine"
	"google.golang.org/appengine/datastore"
)

// Errors
var (
	ErrNoCookie    = errors.New("no cookie found")
	ErrNoAuthToken = errors.New("no header auth token found")
)

// GetToken returns the *Token value for the raw token value contained within the auth cookie or auth header
func GetToken(c context.Context, r *http.Request) (*Token, error) {
	uuid, err := getUUIDFromCookie(r)
	if err != nil && err != http.ErrNoCookie {
		return nil, err
	}
	if err == http.ErrNoCookie {
		uuid = getUUIDFromHeader(r)
	}
	if err != nil {
		return nil, err
	}
	if len(uuid) == 0 {
		return nil, ErrNoAuthToken
	}

	tstore := NewTokenStore()
	return tstore.Get(c, uuid)
}

// Signup creates a user account and links up the credentials. Based on the request type an auth cookie
// or header token will be set with an auth token.
func Signup(c context.Context, w http.ResponseWriter, r *http.Request, creds *Credentials, account *Account) error {
	astore := NewAccountStore()
	tstore := NewTokenStore()

	accountKey, err := astore.Create(c, creds, account)
	if err != nil {
		return errors.New("failed to create account")
	}
	token, err := tstore.Create(c, accountKey)

	if accepts(r, "json") {
		setHeaderToken(w, token.UUID)
	} else {
		setAuthCookieToken(w, token.UUID)
	}

	return nil
}

// Signout deletes the token in the response to the client as well as deletes the token
// in the database to ensure it is no longer usable
func Signout(c context.Context, w http.ResponseWriter, r *http.Request) error {
	var err error
	var uuid string

	if accepts(r, "json") {
		uuid = getUUIDFromHeader(r)
		if err != nil {
			return fmt.Errorf("failed to get token from header: %v", err)
		}
		clearHeader(w)
	} else {
		uuid, err = getUUIDFromCookie(r)
		if err != nil {
			return fmt.Errorf("failed to get token from cookie: %v", err)
		}
		clearCookie(w)
	}

	store := NewTokenStore()
	token, err := store.Get(c, uuid)
	if err != nil {
		return err
	}
	err = datastore.Delete(c, token.Key)

	return err
}

// Authenticate .
func Authenticate(c context.Context, w http.ResponseWriter, r *http.Request, creds *Credentials) (*Token, error) {
	var token *Token
	var err error
	if len(creds.ProviderName) > 0 {
		token, err = doExternalAuth(c, creds, appEngineURLGetter{Ctx: c})
	} else {
		token, err = doInternalAuth(c, creds)
	}
	if err != nil {
		return nil, err
	}

	if accepts(r, "json") {
		setHeaderToken(w, token.UUID)
	} else {
		setAuthCookieToken(w, token.UUID)
	}

	return token, nil
}

func doInternalAuth(c context.Context, creds *Credentials) (*Token, error) {
	accountStore := NewAccountStore()
	tokenStore := NewTokenStore()

	accountKey, err := accountStore.GetAccountKeyByCredentials(c, creds)
	if err != nil {
		return nil, fmt.Errorf("getting account key by credentials: %v", err)
	}

	token, err := tokenStore.Create(c, accountKey)
	if err != nil {
		return nil, err
	}

	return token, nil
}

func doExternalAuth(c context.Context, creds *Credentials, urlGetter urlGetter) (*Token, error) {
	var err error
	tokenStore := NewTokenStore()
	accountStore := NewAccountStore()

	switch creds.ProviderName {
	case "facebook":
		err = fbgraphapi.Authenticate(creds.ProviderToken, creds.ProviderID, urlGetter)
	default:
		return nil, errors.New("unknown auth provider")
	}
	if err != nil {
		return nil, fmt.Errorf("authenticate: %v", err)
	}

	accountKey, err := accountStore.GetAccountKeyByCredentials(c, creds)
	if err != nil {
		return nil, fmt.Errorf("getting account key by credentials: %v", err)
	}

	token, err := tokenStore.Create(c, accountKey)
	if err != nil {
		return nil, err
	}

	return token, nil
}

func clearCookie(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     cookieName,
		Path:     "/",
		Expires:  time.Time{},
		HttpOnly: true,
		Value:    "",
		Secure:   !appengine.IsDevAppServer(),
	})
}

func setAuthCookieToken(w http.ResponseWriter, uuid string) {
	http.SetCookie(w, &http.Cookie{
		Name:     cookieName,
		Path:     "/",
		Expires:  time.Now().Add(time.Hour * 24 * 14),
		HttpOnly: true,
		Secure:   !appengine.IsDevAppServer(),
		Value:    uuid,
	})
}

func setHeaderToken(w http.ResponseWriter, uuid string) {
	w.Header().Set("Authorization", uuid)
}

func getUUIDFromHeader(r *http.Request) string {
	return r.Header.Get("Authorization")[len("token="):]
}

func getUUIDFromCookie(r *http.Request) (string, error) {
	cookie, err := r.Cookie(cookieName)
	if err != nil || len(cookie.Value) == 0 {
		return "", ErrNoCookie
	}

	return cookie.Value, nil
}

func clearHeader(w http.ResponseWriter) {
	w.Header().Del("Authorization")
}

func accepts(r *http.Request, t string) bool {
	return strings.Index(r.Header.Get("Accept"), t) > 0
}

// VerifyReferrer middlware validates the referer header matches the request url's host
func VerifyReferrer(c context.Context, w http.ResponseWriter, r *http.Request) context.Context {
	if r.Method != http.MethodPost {
		return c
	}

	err := func() error {
		referrer, err := url.Parse(r.Header.Get("Referer"))
		if err != nil {
			return err
		}
		if referrer.Host != r.Host {
			return fmt.Errorf("BLOCKED: VerifyReferrer: %v - %v", referrer, r.Host)
		}
		return nil
	}()

	if err != nil {
		cc, cancel := context.WithCancel(c)
		cancel()
		return cc
	}
	return c
}
<|endoftext|>


<commit_before>package ci_bot

import (
	"testing"
)

type item struct {
	addr    string
	cidr    string
	matches bool
}

var testdata = []item{
	item{"192.168.1.67", "192.168.1.0/24", true},
	item{"192.168.1.67", "192.168.1.0/28", false},
	item{"192.168.1.67", "0.0.0.0/0", true},
	item{"not_an_ip", "0.0.0.0/0", false},
	item{"192.168.1.67", "not_a_cidr", false},
	item{"", "0.0.0.0/0", false},
	item{"192.168.1.67", "", false},
}

func TestCIDRMatch(t *testing.T) {
	for _, it := range testdata {
		if IsAddrInCIDR(it.addr, it.cidr) != it.matches {
			t.Fatalf("%s in %s should be %t", it.addr, it.cidr, it.matches)
		}
	}
}
<commit_msg>Simplify test data table<commit_after>package ci_bot

import (
	"testing"
)

type item struct {
	addr    string
	cidr    string
	matches bool
}

var testdata = []item{
	{"192.168.1.67", "192.168.1.0/24", true},
	{"192.168.1.67", "192.168.1.0/28", false},
	{"192.168.1.67", "0.0.0.0/0", true},
	{"not_an_ip", "0.0.0.0/0", false},
	{"192.168.1.67", "not_a_cidr", false},
	{"", "0.0.0.0/0", false},
	{"192.168.1.67", "", false},
}

func TestCIDRMatch(t *testing.T) {
	for _, it := range testdata {
		if IsAddrInCIDR(it.addr, it.cidr) != it.matches {
			t.Fatalf("%s in %s should be %t", it.addr, it.cidr, it.matches)
		}
	}
}
<|endoftext|>


<commit_before>package html2text

import (
	"bytes"
	"golang.org/x/net/html"
	"io"
	"regexp"
	"strings"
)

var (
	spacingRe = regexp.MustCompile(`[ \r\n\t]+`)
	newlineRe = regexp.MustCompile(`\n\n+`)
)

func textify(node *html.Node, buf *bytes.Buffer) error {
	var err error
	var noRecurse bool
	switch node.Type {
	case html.TextNode:
		data := strings.Trim(spacingRe.ReplaceAllString(node.Data, " "), "\r\n \t")
		if len(data) > 0 {
			if err = buf.WriteByte('\n'); err != nil {
				return err
			}
			if _, err = buf.WriteString(data); err != nil {
				return err
			}
			buf.WriteByte('\n')
		}
	case html.ElementNode:
		for _, attr := range node.Attr {
			if attr.Key == "href" {
				if _, err = buf.WriteString(" " + attr.Val); err != nil {
					return err
				}
				noRecurse = true
				break
			}
		}
	}
	if !noRecurse {
		beforeLen := buf.Len()
		for c := node.FirstChild; c != nil; c = c.NextSibling {
			if err = textify(c, buf); err != nil {
				return err
			}
		}
		if afterLen := buf.Len(); beforeLen != afterLen {
			buf.WriteByte('\n')
		}
	}
	return nil
}

func FromReader(reader io.Reader) (string, error) {
	buf := &bytes.Buffer{}
	doc, err := html.Parse(reader)
	if err != nil {
		return "", err
	}
	if err = textify(doc, buf); err != nil {
		return "", err
	}
	text := strings.TrimSpace(newlineRe.ReplaceAllString(strings.Replace(buf.String(), "\n ", "\n", -1), "\n\n"))
	return text, nil
}

func FromString(input string) (string, error) {
	text, err := FromReader(strings.NewReader(input))
	if err != nil {
		return "", err
	}
	return text, nil
}
<commit_msg>Adjusting behavior to remove additional elements,<commit_after>package html2text

// Fork of https://github.com/jaytaylor/html2text
// This does additional processing to omit non-content areas and rewrite;
//    a: tags to markdown style links
//    HX: headers to markdown style with prefixed 'X' number of "#" characters
//    li: tags to markdown style unordered lists (even for ol li's)
// Omission from extraction is done via a combinations of;
//   - tag selection
//   - class selection
//   - aria roles
import (
	"bytes"
	"fmt"
	"golang.org/x/net/html"
	"io"
	"log"
	"os"
	"regexp"
	"strconv"
	"strings"
	"text/tabwriter"
)

const (
	OmitClasses = 1 << iota
	OmitIds
	OmitRoles
)

type state struct {
	node        *html.Node
	buf         *bytes.Buffer
	url         string
	omitClasses bool
	omitIds     bool
	omitRoles   bool
	indent      int
}

// Indent spacing (used for nested li's)
const spacing = 4

var (
	spacingRe      = regexp.MustCompile(`[[:space:]]+`)
	newlineRe      = regexp.MustCompile("\n\n+")
	newlineSpaceRe = regexp.MustCompile("\n[ \r\n\t]+\n")
)

// https://developers.whatwg.org/content-models.html#kinds-of-content
var omitTags = []string{
	// Metadata content
	// https://developers.whatwg.org/content-models.html#metadata-content-0
	"base",
	"link", // @todo: unless itemprop attribute is present
	"meta", // @todo: unless itemprop attribute is present
	"noscript",
	"script",
	"style",
	"template",
	"title",

	// Flow control
	// https://developers.whatwg.org/content-models.html#flow-content-0
	//"a",
	"abbr",
	"address",
	//"area", (if it is a descendant of a map element)
	"article",
	//"aside",
	//"audio",
	//"b",
	//"bdi",
	//"bdo",
	"blockquote",
	//"br",
	"button",
	"canvas",
	"cite",
	"code",
	"data",
	"datalist",
	"del",
	//"details",
	//"dfn",
	"dialog",
	//"div",
	//"dl",
	//"em",
	"embed",
	"fieldset",
	//"figure",
	"footer",
	"form",
	//"h1",
	//"h2",
	//"h3",
	//"h4",
	//"h5",
	//"h6",
	//"header",
	//"hgroup",
	"hr",
	//"i",
	"iframe",
	//"img",
	"input",
	"ins",
	"kbd",
	"keygen",
	//"label",
	"link", //(if the itemprop attribute is present)
	//"main",
	"map",
	//"mark",
	"math",
	"menu",
	"meta", //(if the itemprop attribute is present)
	"meter",
	"nav",
	"noscript",
	"object",
	//"ol",
	"output",
	//"p",
	//"pre",
	"progress",
	//"q",
	//"ruby",
	"s",
	//"samp",
	"script",
	//"section",
	"select",
	//"small",
	//"span",
	//"strong",
	"style", //(if the scoped attribute is present)
	//"sub",
	//"sup",
	"svg",
	//"table",
	"template",
	"textarea",
	"time",
	//"u",
	//"ul",
	//"var",
	"video",
	"wbr",

	// Interactive content (we exclude "a" and "img")
	// https://developers.whatwg.org/content-models.html#interactive-content
	// "a",
	"audio",
	"button",
	"details",
	"embed",
	"iframe",
	// "img",
	"input",
	"keygen",
	//"label",
	"object",
	"select",
	"textarea",
	"sorting",
	"video",

	// embedded content (not included above)
	// https://developers.whatwg.org/content-models.html#embedded-content-0
	// "audio",
	"canvas",
	// "embed",
	// "iframe",
	// "img",
	"math",
	"object",
	"svg",
	"video",
}

// If any class contains a *substring* in this list it will be filtered out.
var omitAttrValues = []string{
	"example",
	"footer",
	"header",
	"menu",
	"modal",
	"nav",
	"overlay",
	"social",
	"tip",
	"tooltip",
}

// Classes to always omit.
var alwaysOmit = []string{
	"hidden",
	"hide",
}

// http://www.w3.org/TR/wai-aria/roles#role_definitions_header
var omitAriaRoles = []string{
	"alert",
	"alertdialog",
	//"application",
	//"article",
	"banner",
	"button",
	"checkbox",
	//"columnheader",
	"combobox",
	"command",
	//"complementary",
	"composite",
	"contentinfo",
	//"definition",
	"dialog",
	"directory",
	//"document",
	"form",
	"grid",
	"gridcell",
	"group",
	//"heading",
	"img",
	"input",
	"landmark",
	//"link",
	//"list",
	"listbox",
	//"listitem",
	"log",
	//"main",
	"marquee",
	"math",
	"menu",
	"menubar",
	"menuitem",
	"menuitemcheckbox",
	"menuitemradio",
	"navigation",
	//"note",
	"option",
	"presentation",
	"progressbar",
	"radio",
	"radiogroup",
	"range",
	//"region",
	"roletype",
	//"row",
	//"rowgroup",
	//"rowheader",
	"scrollbar",
	"search",
	//"section",
	//"sectionhead",
	"select",
	"separator",
	"slider",
	"spinbutton",
	"status",
	//"structure",
	"tab",
	"tablist",
	"tabpanel",
	"textbox",
	"timer",
	"toolbar",
	"tooltip",
	//"tree",
	//"treegrid",
	//"treeitem",
	"widget",
	"window",
}

var sectionTags = []string{
	// Sectioning content;
	// https://developers.whatwg.org/content-models.html#sectioning-content-0
	"article",
	"aside",
	"nav",
	"section",

	// Heading content;
	// https://developers.whatwg.org/content-models.html#heading-content
	"h1",
	"h2",
	"h3",
	"h4",
	"h5",
	"h6",
	"hgroup",
}

// Tags that should have trailing newline.
var breakingTags = append(sectionTags, []string{
	"p",
	"div",
	"br",
	"ul",
	"ol",
}...)

var debug bool
var logger *log.Logger

func init() {
	// Currently if env DEBUG is any non null string we debug.
	env := os.Getenv("DEBUG")
	if env != "" {
		debug = true
		logger = log.New(os.Stderr, "html2text: ", log.LstdFlags|log.Lshortfile)
	}
}

func textify(curState *state) error {
	var curNode *html.Node
	var err error
	var noRecurse bool
	var i int64
	curNode = curState.node

	if stringInSlice(curState.node.DataAtom.String(), breakingTags) {
		if _, err = curState.buf.WriteString("\n\n"); err != nil {
			return err
		}
	}

	if stringInSlice(curState.node.DataAtom.String(), sectionTags) {
		if _, err = curState.buf.WriteString("\n\n"); err != nil {
			return err
		}

		last := curState.node.DataAtom.String()[len(curState.node.DataAtom.String())-1:]
		if i, err = strconv.ParseInt(last, 10, 8); err == nil {
			if _, err = curState.buf.WriteString(strings.Repeat("#", int(i)) + " "); err != nil {
				return err
			}
		}
	}

	switch curState.node.Type {
	case html.TextNode:
		data := strings.Trim(spacingRe.ReplaceAllString(curState.node.Data, " "), "\r\n\t")
		if len(data) > 0 {
			if _, err = curState.buf.WriteString(data); err != nil {
				return err
			}
		}
	case html.ElementNode:
		switch curState.node.DataAtom.String() {
		case "a":
			var text string
			text, err = handleATag(curState)
			if _, err = curState.buf.WriteString(text); err != nil {
				return err
			}
		case "ul", "ol":
			curState.indent++
		case "li":
			var text string
			text, err = handleChildren(curState)
			if strings.TrimSpace(text) != "" {
				if _, err = curState.buf.WriteString("\n" + strings.Repeat(" ", spacing*curState.indent) + "* " + text); err != nil {
					return err
				}
			}
		case "table":
			var text string
			text, err = handleChildren(curState)
			if strings.TrimSpace(text) != "" {
				buf := &bytes.Buffer{}
				w := new(tabwriter.Writer)
				w.Init(buf, 0, 4, 2, ' ', tabwriter.Debug)
				fmt.Fprint(w, "|"+strings.Replace(text, "\n", "\n|", -1))
				w.Flush()
				if _, err = curState.buf.WriteString("\n\n" + buf.String() + "\n\n"); err != nil {
					return err
				}
			}
		case "td", "th":
			var text string
			text, err = handleChildren(curState)
			if strings.TrimSpace(text) != "" {
				data := strings.TrimSpace(spacingRe.ReplaceAllString(text, " "))
				if len(data) > 0 {
					if _, err = curState.buf.WriteString(data); err != nil {
						return err
					}
				}
			} else {

			}
		default:
			// Stop checking at first noRecurse.

			if curState.node.DataAtom.String() != "" && stringInSlice(curState.node.DataAtom.String(), omitTags) {
				noRecurse = true
				break
			}

			if AttrHasString(curState.node, "class", alwaysOmit) && curNode.DataAtom.String() != "body" {
				if debug {
					logger.Printf("[%s] skipping recursion for %s bc class\n", curState.url, curState.node.DataAtom.String())
				}
				noRecurse = true
				break
			}

			// Omit classes except when they appear on the body
			if curState.omitClasses {
				if AttrHasString(curState.node, "class", omitAttrValues) && curNode.DataAtom.String() != "body" {
					if debug {
						logger.Printf("[%s] skipping recursion for %s bc class\n", curState.url, curState.node.DataAtom.String())
					}
					noRecurse = true
					break
				}
			}

			// Omit ids except when they appear on the body
			if curState.omitIds {
				if AttrHasString(curState.node, "id", omitAttrValues) && curNode.DataAtom.String() != "body" {
					if debug {
						logger.Printf("[%s] skipping recursion for %s bc id\n", curState.url, curState.node.DataAtom.String())
					}
					noRecurse = true
					break
				}
			}

			if curState.omitRoles {
				if AttrHasString(curState.node, "role", omitAriaRoles) {
					if debug {
						logger.Printf("[%s] skipping recursion for %s bc role\n", curState.url, curState.node.DataAtom.String())
					}
					noRecurse = true
					break
				}
			}

		}
	}
	if !noRecurse {
		for c := curState.node.FirstChild; c != nil; c = c.NextSibling {
			curState.node = c
			if err = textify(curState); err != nil {
				return err
			}
		}
	}

	// After tag work.
	switch curNode.DataAtom.String() {
	case "ul", "ol":
		curState.indent--
	case "td", "th":
		if err = curState.buf.WriteByte('\t'); err != nil {
			return err
		}
	case "tr":
		if err = curState.buf.WriteByte('\n'); err != nil {
			return err
		}
	}

	if stringInSlice(curState.node.DataAtom.String(), breakingTags) {
		if _, err = curState.buf.WriteString("\n\n"); err != nil {
			return err
		}
	}

	return nil
}

// Recurse down an "a" node tree and collect all node content and merge to
// one line of text. This is so that nested divs and other elements don't
// create excessive whitespace.
func handleATag(curState *state) (string, error) {
	var text string
	var err error
	for _, attr := range curState.node.Attr {
		if attr.Key == "href" {
			// Relative -> Absolute URL
			var link string
			if strings.HasPrefix(attr.Val, "#") || strings.HasPrefix(attr.Val, curState.url+"#") {
				// We don't want to capture internal links
			} else if strings.HasPrefix(attr.Val, "/") {
				link = curState.url + attr.Val
			} else if strings.HasPrefix(attr.Val, ".") {
				link = curState.url + "/" + attr.Val
			} else {
				link = attr.Val
			}
			buf := curState.buf
			curState.buf = &bytes.Buffer{}
			for c := curState.node.FirstChild; c != nil; c = c.NextSibling {
				curState.node = c
				if err = textify(curState); err != nil {
					return text, err
				}
			}
			if link != "" {
				linkText := strings.Trim(spacingRe.ReplaceAllString(curState.buf.String(), " "), "\r\n \t")
				text = " (" + link + ")[" + linkText + "]"
			}
			curState.buf = buf
			break
		}
	}
	return text, err
}

// Accumulates child content.
func handleChildren(curState *state) (string, error) {
	var text string
	var err error

	buf := curState.buf
	curState.buf = &bytes.Buffer{}
	for c := curState.node.FirstChild; c != nil; c = c.NextSibling {
		curState.node = c
		if err = textify(curState); err != nil {
			return text, err
		}
	}
	text = curState.buf.String()
	curState.buf = buf

	return text, err
}

func FromReader(reader io.Reader, url string, flags int) (string, error) {
	var err error
	curState := &state{}
	curState.buf = &bytes.Buffer{}
	curState.url = url
	if (flags & OmitClasses) != 0 {
		curState.omitClasses = true
	}
	if (flags & OmitIds) != 0 {
		curState.omitIds = true
	}
	if (flags & OmitRoles) != 0 {
		curState.omitRoles = true
	}
	curState.node, err = html.Parse(reader)
	if err != nil {
		return "", err
	}
	if err = textify(curState); err != nil {
		return "", err
	}
	text := strings.TrimSpace(newlineSpaceRe.ReplaceAllString(curState.buf.String(), "\n\n"))
	return text, nil
}

func stringInSlice(s string, list []string) bool {
	for _, b := range list {
		if strings.Contains(s, b) {
			return true
		}
	}
	return false
}

func AttrHasString(n *html.Node, key string, values []string) bool {
	for _, attr := range n.Attr {
		if attr.Key == key {
			if stringInSlice(attr.Val, values) {
				if debug {
					logger.Printf("node(%s) has '%s' with '%s' matching string in (%s)\n", n.DataAtom.String(), key, attr.Val, strings.Join(values, ", "))
				}
				return true
			}
			break
		}
	}
	return false
}

func FromString(input string, url string, flags int) (string, error) {
	text, err := FromReader(strings.NewReader(input), url, flags)
	if err != nil {
		return "", err
	}
	return text, nil
}
<|endoftext|>


<commit_before>package netallow

import (
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
)

type testHandler struct {
	Message string
}

func newTestHandler(m string) http.Handler {
	return &testHandler{Message: m}
}

func (h *testHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte(h.Message))
}

var testAllowHandler = newTestHandler("OK")
var testDenyHandler = newTestHandler("NO")

func testHTTPResponse(url string, t *testing.T) string {
	resp, err := http.Get(url)
	if err != nil {
		t.Fatalf("%v", err)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("%v", err)
	}
	resp.Body.Close()
	return string(body)
}

func testWorker(url string, t *testing.T, wg *sync.WaitGroup) {
	for i := 0; i < 100; i++ {
		response := testHTTPResponse(url, t)
		if response != "NO" {
			t.Fatalf("Expected NO, but got %s", response)
		}
	}
	wg.Done()
}

func TestHostStubHTTP(t *testing.T) {
	acl := NewHostStub()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	response := testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}

	addIPString(acl, "127.0.0.1", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}

	delIPString(acl, "127.0.0.1", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}
}

func TestNetStubHTTP(t *testing.T) {
	acl := NewNetStub()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	response := testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}

	testAddNet(acl, "127.0.0.1/32", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}

	testDelNet(acl, "127.0.0.1/32", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}
}

func TestBasicHTTP(t *testing.T) {
	acl := NewBasic()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	response := testHTTPResponse(srv.URL, t)
	if response != "NO" {
		t.Fatalf("Expected NO, but got %s", response)
	}

	addIPString(acl, "127.0.0.1", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}

	delIPString(acl, "127.0.0.1", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "NO" {
		t.Fatalf("Expected NO, but got %s", response)
	}
}

func TestBasicHTTPDefaultDeny(t *testing.T) {
	acl := NewBasic()
	h, err := NewHandler(testAllowHandler, nil, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	expected := "Unauthorized"
	response := strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}
}

func TestBasicHTTPWorkers(t *testing.T) {
	acl := NewBasic()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	wg := new(sync.WaitGroup)
	defer srv.Close()

	for i := 0; i < 16; i++ {
		wg.Add(1)
		go testWorker(srv.URL, t, wg)
	}

	wg.Wait()

}

func TestFailHTTP(t *testing.T) {
	acl := NewBasic()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	w := httptest.NewRecorder()
	req := new(http.Request)

	if h.ServeHTTP(w, req); w.Code != http.StatusInternalServerError {
		t.Fatalf("Expect HTTP 500, but got HTTP %d", w.Code)
	}
}

var testHandlerFunc *HandlerFunc

func newTestHandlerFunc(m string) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(m))
	}
}

var testAllowHandlerFunc = newTestHandlerFunc("OK")
var testDenyHandlerFunc = newTestHandlerFunc("NO")

func TestSetupHandlerFuncFails(t *testing.T) {
	acl := NewBasic()
	_, err := NewHandlerFunc(nil, testDenyHandlerFunc, acl)
	if err == nil {
		t.Fatal("expected NewHandlerFunc to fail with nil allow handler")
	}

	_, err = NewHandlerFunc(testAllowHandlerFunc, testDenyHandlerFunc, nil)
	if err == nil {
		t.Fatal("expected NewHandlerFunc to fail with nil allowed")
	}

	_, err = NewHandlerFunc(testAllowHandlerFunc, nil, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}
}

func TestSetupHandlerFunc(t *testing.T) {
	acl := NewBasic()
	h, err := NewHandlerFunc(testAllowHandlerFunc, testDenyHandlerFunc, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	expected := "NO"
	response := strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}

	h.deny = nil
	expected = "Unauthorized"
	response = strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}

	addIPString(acl, "127.0.0.1", t)
	expected = "OK"
	response = strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}
}

func TestFailHTTPFunc(t *testing.T) {
	acl := NewBasic()
	h, err := NewHandlerFunc(testAllowHandlerFunc, testDenyHandlerFunc, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	w := httptest.NewRecorder()
	req := new(http.Request)

	if h.ServeHTTP(w, req); w.Code != http.StatusInternalServerError {
		t.Fatalf("Expect HTTP 500, but got HTTP %d", w.Code)
	}
}

func TestBasicNetHTTP(t *testing.T) {
	acl := NewBasicNet()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	response := testHTTPResponse(srv.URL, t)
	if response != "NO" {
		t.Fatalf("Expected NO, but got %s", response)
	}

	testAddNet(acl, "127.0.0.1/32", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}

	testDelNet(acl, "127.0.0.1/32", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "NO" {
		t.Fatalf("Expected NO, but got %s", response)
	}
}

func TestBasicNetHTTPDefaultDeny(t *testing.T) {
	acl := NewBasicNet()
	h, err := NewHandler(testAllowHandler, nil, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	expected := "Unauthorized"
	response := strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}
}

func TestBasicNetHTTPWorkers(t *testing.T) {
	acl := NewBasicNet()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	wg := new(sync.WaitGroup)
	defer srv.Close()

	for i := 0; i < 16; i++ {
		wg.Add(1)
		go testWorker(srv.URL, t, wg)
	}

	wg.Wait()

}

func TestNetFailHTTP(t *testing.T) {
	acl := NewBasicNet()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}
	w := httptest.NewRecorder()
	req := new(http.Request)

	if h.ServeHTTP(w, req); w.Code != http.StatusInternalServerError {
		t.Fatalf("Expect HTTP 500, but got HTTP %d", w.Code)
	}
}

func TestSetupNetHandlerFuncFails(t *testing.T) {
	acl := NewBasicNet()
	_, err := NewHandlerFunc(nil, testDenyHandlerFunc, acl)
	if err == nil {
		t.Fatal("expected NewHandlerFunc to fail with nil allow handler")
	}

	_, err = NewHandlerFunc(testAllowHandlerFunc, testDenyHandlerFunc, nil)
	if err == nil {
		t.Fatal("expected NewHandlerFunc to fail with nil allowed")
	}

	_, err = NewHandlerFunc(testAllowHandlerFunc, nil, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}
}

func TestSetupNetHandlerFunc(t *testing.T) {
	acl := NewBasicNet()
	h, err := NewHandlerFunc(testAllowHandlerFunc, testDenyHandlerFunc, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	expected := "NO"
	response := strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}

	h.deny = nil
	expected = "Unauthorized"
	response = strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}

	testAddNet(acl, "127.0.0.1/32", t)
	expected = "OK"
	response = strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}
}

func TestNetFailHTTPFunc(t *testing.T) {
	acl := NewBasicNet()
	h, err := NewHandlerFunc(testAllowHandlerFunc, testDenyHandlerFunc, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	w := httptest.NewRecorder()
	req := new(http.Request)

	if h.ServeHTTP(w, req); w.Code != http.StatusInternalServerError {
		t.Fatalf("Expect HTTP 500, but got HTTP %d", w.Code)
	}
}

func TestHandlerFunc(t *testing.T) {
	var acl ACL
	_, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err == nil || err.Error() != "allowed: ACL cannot be nil" {
		t.Fatal("Expected error with nil allow handler.")
	}

	acl = NewBasic()
	_, err = NewHandler(nil, testDenyHandler, acl)
	if err == nil || err.Error() != "allowed: allow cannot be nil" {
		t.Fatal("Expected error with nil ACL.")
	}
}
<commit_msg>Third time's the charm?<commit_after>package netallow

import (
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
)

type testHandler struct {
	Message string
}

func newTestHandler(m string) http.Handler {
	return &testHandler{Message: m}
}

func (h *testHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte(h.Message))
}

var testAllowHandler = newTestHandler("OK")
var testDenyHandler = newTestHandler("NO")

func testHTTPResponse(url string, t *testing.T) string {
	resp, err := http.Get(url)
	if err != nil {
		t.Fatalf("%v", err)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("%v", err)
	}
	resp.Body.Close()
	return string(body)
}

func testWorker(url string, t *testing.T, wg *sync.WaitGroup) {
	for i := 0; i < 100; i++ {
		response := testHTTPResponse(url, t)
		if response != "NO" {
			t.Fatalf("Expected NO, but got %s", response)
		}
	}
	wg.Done()
}

func TestHostStubHTTP(t *testing.T) {
	acl := NewHostStub()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	response := testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}

	addIPString(acl, "127.0.0.1", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}

	delIPString(acl, "127.0.0.1", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}
}

func TestNetStubHTTP(t *testing.T) {
	acl := NewNetStub()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	response := testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}

	testAddNet(acl, "127.0.0.1/32", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}

	testDelNet(acl, "127.0.0.1/32", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}
}

func TestBasicHTTP(t *testing.T) {
	acl := NewBasic()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	response := testHTTPResponse(srv.URL, t)
	if response != "NO" {
		t.Fatalf("Expected NO, but got %s", response)
	}

	addIPString(acl, "127.0.0.1", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}

	delIPString(acl, "127.0.0.1", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "NO" {
		t.Fatalf("Expected NO, but got %s", response)
	}
}

func TestBasicHTTPDefaultDeny(t *testing.T) {
	acl := NewBasic()
	h, err := NewHandler(testAllowHandler, nil, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	expected := "Unauthorized"
	response := strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}
}

func TestBasicHTTPWorkers(t *testing.T) {
	acl := NewBasic()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	wg := new(sync.WaitGroup)
	defer srv.Close()

	for i := 0; i < 16; i++ {
		wg.Add(1)
		go testWorker(srv.URL, t, wg)
	}

	wg.Wait()

}

func TestFailHTTP(t *testing.T) {
	acl := NewBasic()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	w := httptest.NewRecorder()
	req := new(http.Request)

	if h.ServeHTTP(w, req); w.Code != http.StatusInternalServerError {
		t.Fatalf("Expect HTTP 500, but got HTTP %d", w.Code)
	}
}

var testHandlerFunc *HandlerFunc

func newTestHandlerFunc(m string) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(m))
	}
}

var testAllowHandlerFunc = newTestHandlerFunc("OK")
var testDenyHandlerFunc = newTestHandlerFunc("NO")

func TestSetupHandlerFuncFails(t *testing.T) {
	acl := NewBasic()
	_, err := NewHandlerFunc(nil, testDenyHandlerFunc, acl)
	if err == nil {
		t.Fatal("expected NewHandlerFunc to fail with nil allow handler")
	}

	_, err = NewHandlerFunc(testAllowHandlerFunc, testDenyHandlerFunc, nil)
	if err == nil {
		t.Fatal("expected NewHandlerFunc to fail with nil allowed")
	}

	_, err = NewHandlerFunc(testAllowHandlerFunc, nil, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}
}

func TestSetupHandlerFunc(t *testing.T) {
	acl := NewBasic()
	h, err := NewHandlerFunc(testAllowHandlerFunc, testDenyHandlerFunc, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	expected := "NO"
	response := strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}

	h.deny = nil
	expected = "Unauthorized"
	response = strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}

	addIPString(acl, "127.0.0.1", t)
	expected = "OK"
	response = strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}
}

func TestFailHTTPFunc(t *testing.T) {
	acl := NewBasic()
	h, err := NewHandlerFunc(testAllowHandlerFunc, testDenyHandlerFunc, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	w := httptest.NewRecorder()
	req := new(http.Request)

	if h.ServeHTTP(w, req); w.Code != http.StatusInternalServerError {
		t.Fatalf("Expect HTTP 500, but got HTTP %d", w.Code)
	}
}

func TestBasicNetHTTP(t *testing.T) {
	acl := NewBasicNet()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	response := testHTTPResponse(srv.URL, t)
	if response != "NO" {
		t.Fatalf("Expected NO, but got %s", response)
	}

	testAddNet(acl, "127.0.0.1/32", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "OK" {
		t.Fatalf("Expected OK, but got %s", response)
	}

	testDelNet(acl, "127.0.0.1/32", t)
	response = testHTTPResponse(srv.URL, t)
	if response != "NO" {
		t.Fatalf("Expected NO, but got %s", response)
	}
}

func TestBasicNetHTTPDefaultDeny(t *testing.T) {
	acl := NewBasicNet()
	h, err := NewHandler(testAllowHandler, nil, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	expected := "Unauthorized"
	response := strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}
}

func TestBasicNetHTTPWorkers(t *testing.T) {
	acl := NewBasicNet()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	wg := new(sync.WaitGroup)
	defer srv.Close()

	for i := 0; i < 16; i++ {
		wg.Add(1)
		go testWorker(srv.URL, t, wg)
	}

	wg.Wait()

}

func TestNetFailHTTP(t *testing.T) {
	acl := NewBasicNet()
	h, err := NewHandler(testAllowHandler, testDenyHandler, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}
	w := httptest.NewRecorder()
	req := new(http.Request)

	if h.ServeHTTP(w, req); w.Code != http.StatusInternalServerError {
		t.Fatalf("Expect HTTP 500, but got HTTP %d", w.Code)
	}
}

func TestSetupNetHandlerFuncFails(t *testing.T) {
	acl := NewBasicNet()
	_, err := NewHandlerFunc(nil, testDenyHandlerFunc, acl)
	if err == nil {
		t.Fatal("expected NewHandlerFunc to fail with nil allow handler")
	}

	_, err = NewHandlerFunc(testAllowHandlerFunc, testDenyHandlerFunc, nil)
	if err == nil {
		t.Fatal("expected NewHandlerFunc to fail with nil allowed")
	}

	_, err = NewHandlerFunc(testAllowHandlerFunc, nil, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}
}

func TestSetupNetHandlerFunc(t *testing.T) {
	acl := NewBasicNet()
	h, err := NewHandlerFunc(testAllowHandlerFunc, testDenyHandlerFunc, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	srv := httptest.NewServer(h)
	defer srv.Close()

	expected := "NO"
	response := strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}

	h.deny = nil
	expected = "Unauthorized"
	response = strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}

	testAddNet(acl, "127.0.0.1/32", t)
	expected = "OK"
	response = strings.TrimSpace(testHTTPResponse(srv.URL, t))
	if response != expected {
		t.Fatalf("Expected %s, but got %s", expected, response)
	}
}

func TestNetFailHTTPFunc(t *testing.T) {
	acl := NewBasicNet()
	h, err := NewHandlerFunc(testAllowHandlerFunc, testDenyHandlerFunc, acl)
	if err != nil {
		t.Fatalf("%v", err)
	}

	w := httptest.NewRecorder()
	req := new(http.Request)

	if h.ServeHTTP(w, req); w.Code != http.StatusInternalServerError {
		t.Fatalf("Expect HTTP 500, but got HTTP %d", w.Code)
	}
}

func TestHandlerFunc(t *testing.T) {
	var acl ACL
	_, err := NewHandler(testAllowHandler, testDenyHandler, nil)
	if err == nil || err.Error() != "netallow: ACL cannot be nil" {
		t.Fatal("Expected error with nil ACL.")
	}

	acl = NewBasic()
	_, err = NewHandler(nil, testDenyHandler, acl)
	if err == nil || err.Error() != "netallow: allow cannot be nil" {
		t.Fatal("Expected error with nil allow handler.")
	}
}
<|endoftext|>


<commit_before>package word2vec

import (
	"errors"
	"fmt"
	"net/http/httptest"
	"reflect"
	"strings"
	"testing"
)

type testCoser struct {
	cos   func(x, y Expr) (float32, error)
	coses func(pairs [][2]Expr) ([]float32, error)
	cosN  func(x Expr, n int) ([]Match, error)
}

func (t testCoser) Cos(x, y Expr) (float32, error)           { return t.cos(x, y) }
func (t testCoser) Coses(pairs [][2]Expr) ([]float32, error) { return t.coses(pairs) }
func (t testCoser) CosN(x Expr, n int) ([]Match, error)      { return t.cosN(x, n) }

func TestEndToEndCos(t *testing.T) {
	tc := &testCoser{}
	h := NewServer(tc)
	s := httptest.NewServer(h)
	defer s.Close()

	c := Client{
		Addr: strings.TrimPrefix(s.URL, "http://"),
	}

	cosTests := []struct {
		x, y Expr
		f    float32
		err  error
	}{
		{
			x:   Expr{"hello": 1.0},
			y:   Expr{"world": 1.0},
			f:   1.0,
			err: nil,
		},
		{
			x:   Expr{"": 1.0},
			y:   Expr{"empty": 1.0},
			f:   0,
			err: fmt.Errorf("suffix test"),
		},
		{
			x:   Expr{"hello": 1.0, "something": 2.0},
			y:   Expr{"world": 1.0, "else": 0.0},
			f:   1.0,
			err: nil,
		},
	}

	for _, tt := range cosTests {
		var cosX, cosY Expr

		tc.cos = func(x, y Expr) (float32, error) {
			cosX = x
			cosY = y
			return tt.f, tt.err
		}

		f, err := c.Cos(tt.x, tt.y)
		if !reflect.DeepEqual(cosX, tt.x) {
			t.Errorf("cosX = %#v, expected: %#v", cosX, tt.x)
		}
		if !reflect.DeepEqual(cosY, tt.y) {
			t.Errorf("cosY = %#v, expected: %#v", cosY, tt.y)
		}
		if f != tt.f {
			t.Errorf("f = %#v, expected: %#v", f, tt.f)
		}
		if tt.err != nil {
			if err == nil {
				t.Errorf("err = %#v, expected: %#v", err, tt.err)
			}
			if !strings.HasSuffix(err.Error(), tt.err.Error()) {
				t.Errorf("err = %q, expected suffix: %q", err, tt.err)
			}
		}
	}
}

func TestEndToEndCoses(t *testing.T) {
	tc := &testCoser{}
	h := NewServer(tc)
	s := httptest.NewServer(h)
	defer s.Close()

	c := Client{
		Addr: strings.TrimPrefix(s.URL, "http://"),
	}

	cosesTests := []struct {
		pairs [][2]Expr
		f     []float32
		err   error
	}{
		{
			pairs: [][2]Expr{
				{{"hello": 1.0}, {"world": 1.0}},
			},
			f:   []float32{1.0},
			err: nil,
		},
		{
			pairs: [][2]Expr{
				{{"hello": 1.0}, {"world": 1.0}},
			},
			f:   nil,
			err: errors.New("suffix"),
		},
	}

	for _, tt := range cosesTests {
		var cosesPairs [][2]Expr
		tc.coses = func(pairs [][2]Expr) ([]float32, error) {
			cosesPairs = pairs
			return tt.f, tt.err
		}

		f, err := c.Coses(tt.pairs)
		if !reflect.DeepEqual(cosesPairs, tt.pairs) {
			t.Errorf("cosesPairs = %#v, expected: %#v", cosesPairs, tt.pairs)
		}
		if !reflect.DeepEqual(f, tt.f) {
			t.Errorf("f = %#v, expected: %#v", f, tt.f)
		}
		if tt.err != nil {
			if err == nil {
				t.Errorf("err = %#v, expected: %#v", err, tt.err)
			}
			if !strings.HasSuffix(err.Error(), tt.err.Error()) {
				t.Errorf("err = %q, expected suffix: %q", err, tt.err)
			}
		}
	}
}

func TestEndToEndCosN(t *testing.T) {
	tc := &testCoser{}
	h := NewServer(tc)
	s := httptest.NewServer(h)
	defer s.Close()

	c := Client{
		Addr: strings.TrimPrefix(s.URL, "http://"),
	}

	cosNTests := []struct {
		x   Expr
		n   int
		m   []Match
		err error
	}{
		{
			x:   Expr{"hello": 1.0},
			n:   0,
			m:   []Match{},
			err: nil,
		},
		{
			x: Expr{"hello": 1.0},
			n: 10,
			m: []Match{
				{"hello", 1.0},
				{"something", 2.0},
				{"else", 1.0},
				{"is", 0.5},
				{"in", 0.1},
				{"the", 0.4},
			},
			err: nil,
		},
	}

	for _, tt := range cosNTests {
		var cosNN int
		var cosNX Expr
		tc.cosN = func(x Expr, n int) ([]Match, error) {
			cosNX = x
			cosNN = n
			return tt.m, tt.err
		}

		m, err := c.CosN(tt.x, tt.n)
		if !reflect.DeepEqual(cosNX, tt.x) {
			t.Errorf("cosNX = %#v, expected: %#v", cosNX, tt.x)
		}
		if !reflect.DeepEqual(cosNN, tt.n) {
			t.Errorf("cosNN = %#v, expected: %#v", cosNN, tt.n)
		}
		if !reflect.DeepEqual(m, tt.m) {
			t.Errorf("m = %#v, expected: %#v", m, tt.m)
		}
		if tt.err != nil {
			if err == nil {
				t.Errorf("err = %#v, expected: %#v", err, tt.err)
			}
			if !strings.HasSuffix(err.Error(), tt.err.Error()) {
				t.Errorf("err = %q, expected suffix: %q", err, tt.err)
			}
		}
	}
}
<commit_msg>Added test for CosN returning error.<commit_after>package word2vec

import (
	"errors"
	"fmt"
	"net/http/httptest"
	"reflect"
	"strings"
	"testing"
)

type testCoser struct {
	cos   func(x, y Expr) (float32, error)
	coses func(pairs [][2]Expr) ([]float32, error)
	cosN  func(x Expr, n int) ([]Match, error)
}

func (t testCoser) Cos(x, y Expr) (float32, error)           { return t.cos(x, y) }
func (t testCoser) Coses(pairs [][2]Expr) ([]float32, error) { return t.coses(pairs) }
func (t testCoser) CosN(x Expr, n int) ([]Match, error)      { return t.cosN(x, n) }

func TestEndToEndCos(t *testing.T) {
	tc := &testCoser{}
	h := NewServer(tc)
	s := httptest.NewServer(h)
	defer s.Close()

	c := Client{
		Addr: strings.TrimPrefix(s.URL, "http://"),
	}

	cosTests := []struct {
		x, y Expr
		f    float32
		err  error
	}{
		{
			x:   Expr{"hello": 1.0},
			y:   Expr{"world": 1.0},
			f:   1.0,
			err: nil,
		},
		{
			x:   Expr{"": 1.0},
			y:   Expr{"empty": 1.0},
			f:   0,
			err: fmt.Errorf("suffix test"),
		},
		{
			x:   Expr{"hello": 1.0, "something": 2.0},
			y:   Expr{"world": 1.0, "else": 0.0},
			f:   1.0,
			err: nil,
		},
	}

	for _, tt := range cosTests {
		var cosX, cosY Expr

		tc.cos = func(x, y Expr) (float32, error) {
			cosX = x
			cosY = y
			return tt.f, tt.err
		}

		f, err := c.Cos(tt.x, tt.y)
		if !reflect.DeepEqual(cosX, tt.x) {
			t.Errorf("cosX = %#v, expected: %#v", cosX, tt.x)
		}
		if !reflect.DeepEqual(cosY, tt.y) {
			t.Errorf("cosY = %#v, expected: %#v", cosY, tt.y)
		}
		if f != tt.f {
			t.Errorf("f = %#v, expected: %#v", f, tt.f)
		}
		if tt.err != nil {
			if err == nil {
				t.Errorf("err = %#v, expected: %#v", err, tt.err)
			}
			if !strings.HasSuffix(err.Error(), tt.err.Error()) {
				t.Errorf("err = %q, expected suffix: %q", err, tt.err)
			}
		}
	}
}

func TestEndToEndCoses(t *testing.T) {
	tc := &testCoser{}
	h := NewServer(tc)
	s := httptest.NewServer(h)
	defer s.Close()

	c := Client{
		Addr: strings.TrimPrefix(s.URL, "http://"),
	}

	cosesTests := []struct {
		pairs [][2]Expr
		f     []float32
		err   error
	}{
		{
			pairs: [][2]Expr{
				{{"hello": 1.0}, {"world": 1.0}},
			},
			f:   []float32{1.0},
			err: nil,
		},
		{
			pairs: [][2]Expr{
				{{"hello": 1.0}, {"world": 1.0}},
			},
			f:   nil,
			err: errors.New("suffix"),
		},
	}

	for _, tt := range cosesTests {
		var cosesPairs [][2]Expr
		tc.coses = func(pairs [][2]Expr) ([]float32, error) {
			cosesPairs = pairs
			return tt.f, tt.err
		}

		f, err := c.Coses(tt.pairs)
		if !reflect.DeepEqual(cosesPairs, tt.pairs) {
			t.Errorf("cosesPairs = %#v, expected: %#v", cosesPairs, tt.pairs)
		}
		if !reflect.DeepEqual(f, tt.f) {
			t.Errorf("f = %#v, expected: %#v", f, tt.f)
		}
		if tt.err != nil {
			if err == nil {
				t.Errorf("err = %#v, expected: %#v", err, tt.err)
			}
			if !strings.HasSuffix(err.Error(), tt.err.Error()) {
				t.Errorf("err = %q, expected suffix: %q", err, tt.err)
			}
		}
	}
}

func TestEndToEndCosN(t *testing.T) {
	tc := &testCoser{}
	h := NewServer(tc)
	s := httptest.NewServer(h)
	defer s.Close()

	c := Client{
		Addr: strings.TrimPrefix(s.URL, "http://"),
	}

	cosNTests := []struct {
		x   Expr
		n   int
		m   []Match
		err error
	}{
		{
			x:   Expr{"hello": 1.0},
			n:   0,
			m:   []Match{},
			err: nil,
		},
		{
			x: Expr{"hello": 1.0},
			n: 10,
			m: []Match{
				{"hello", 1.0},
				{"something", 2.0},
				{"else", 1.0},
				{"is", 0.5},
				{"in", 0.1},
				{"the", 0.4},
			},
			err: nil,
		},
		{
			x:   Expr{"hello": 1.0},
			n:   10,
			m:   nil,
			err: fmt.Errorf("suffix error"),
		},
	}

	for _, tt := range cosNTests {
		var cosNN int
		var cosNX Expr
		tc.cosN = func(x Expr, n int) ([]Match, error) {
			cosNX = x
			cosNN = n
			return tt.m, tt.err
		}

		m, err := c.CosN(tt.x, tt.n)
		if !reflect.DeepEqual(cosNX, tt.x) {
			t.Errorf("cosNX = %#v, expected: %#v", cosNX, tt.x)
		}
		if !reflect.DeepEqual(cosNN, tt.n) {
			t.Errorf("cosNN = %#v, expected: %#v", cosNN, tt.n)
		}
		if !reflect.DeepEqual(m, tt.m) {
			t.Errorf("m = %#v, expected: %#v", m, tt.m)
		}
		if tt.err != nil {
			if err == nil {
				t.Errorf("err = %#v, expected: %#v", err, tt.err)
			}
			if !strings.HasSuffix(err.Error(), tt.err.Error()) {
				t.Errorf("err = %q, expected suffix: %q", err, tt.err)
			}
		}
	}
}
<|endoftext|>