<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	_ "github.com/bmizerany/pq"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

const OPEN311_API_URI = "http://311api.cityofchicago.org/open311/v2/requests.json?extensions=true&page_size=500"

type Open311Request struct {
	Lat, Long                                                                                float64
	Ward, Police_district                                                                                int
	Service_request_id, Status, Service_name, Service_code, Agency_responsible, Address, Channel, Media_url string
	Requested_datetime, Updated_datetime                                                                    string // FIXME: should these be proper time objects?
	Extended_attributes                                                                                map[string]interface{}
}

func main() {
	// open database
	db, err := sql.Open("postgres", "dbname=cwfy sslmode=disable")
	if err != nil {
		log.Fatal("Cannot open database connection", err)
	}
	defer db.Close()

	// load requests from open311
	requests := fetchRequests()

	for _, request := range requests {
		// for each request, either create or update the
		// corresponding record in the database.

		if request.Service_request_id == "" {
			log.Printf("Ignoring a request type %s because there is no SR number assigned", request.Service_name)
			continue
		}

		insert_stmt, err := db.Prepare("INSERT INTO service_requests(service_request_id," +
			"status, service_name, service_code, agency_responsible, " +
			"address, requested_datetime, updated_datetime, lat, long," +
			"ward, police_district, media_url, channel, duplicate, parent_service_request_id) " +
			"SELECT $1::varchar, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16 " +
			"WHERE NOT EXISTS (SELECT 1 FROM service_requests WHERE service_request_id = $1);")

		if err != nil {
			log.Fatal("error preparing database insert statement", err)
		}

		update_stmt, err := db.Prepare("UPDATE service_requests SET " +
			"status = $2, service_name = $3, service_code = $4, agency_responsible = $5, " +
			"address = $6, requested_datetime = $7, updated_datetime = $8, lat = $9, long = $10," +
			"ward = $11, police_district = $12, media_url = $13, channel = $14, duplicate = $15, " +
			"parent_service_request_id = $16, updated_at = NOW() WHERE service_request_id = $1;")

		if err != nil {
			log.Fatal("error preparing database update statement", err)
		}

		tx, err := db.Begin()

		if err != nil {
			log.Fatal("error beginning transaction", err)
		}

		_, err = tx.Stmt(update_stmt).Exec(request.Service_request_id,
			request.Status,
			request.Service_name,
			request.Service_code,
			request.Agency_responsible,
			request.Address,
			request.Requested_datetime,
			request.Updated_datetime,
			request.Lat,
			request.Long,
			request.Extended_attributes["ward"],
			request.Extended_attributes["police_district"],
			request.Media_url,
			request.Extended_attributes["channel"],
			request.Extended_attributes["duplicate"],
			request.Extended_attributes["parent_service_request_id"])

		if err != nil {
			log.Fatalf("could not update %s because %s", request.Service_request_id, err)
		}

		_, err = tx.Stmt(insert_stmt).Exec(request.Service_request_id,
			request.Status,
			request.Service_name,
			request.Service_code,
			request.Agency_responsible,
			request.Address,
			request.Requested_datetime,
			request.Updated_datetime,
			request.Lat,
			request.Long,
			request.Extended_attributes["ward"],
			request.Extended_attributes["police_district"],
			request.Media_url,
			request.Extended_attributes["channel"],
			request.Extended_attributes["duplicate"],
			request.Extended_attributes["parent_service_request_id"])

		if err != nil {
			log.Fatalf("could not save %s because %s", request.Service_request_id, err)
		} else {
			log.Printf("saved SR %s", request)
		}

		err = tx.Commit()
		if err != nil {
			log.Fatal("error closing transaction", err)
		}
	}
}

func (req Open311Request) String() string {
	// pretty print SR information
	return fmt.Sprintf("%s: %s at %s %f,%f, last update %s", req.Service_request_id, req.Service_name, req.Address, req.Lat, req.Long, req.Updated_datetime)
}

func fetchRequests() (requests []Open311Request) {
	db, err := sql.Open("postgres", "dbname=cwfy sslmode=disable")
	if err != nil {
		log.Fatal("Cannot open database connection", err)
	}
	defer db.Close()

	// find the most recent SR that we know about in the database
	rows, err := db.Query("SELECT MAX(updated_datetime) FROM service_requests;")
	if err != nil {
		log.Fatal("error finding most recent service request", err)
	}

	last_updated_at := time.Now()
	for rows.Next() {
		if err := rows.Scan(&last_updated_at); err != nil {
			log.Fatal(err)
		}

		log.Printf("most recent SR timestamp %s", last_updated_at)
	}

	// janky hack to transform the last updated timestamp into
	// a format that plays nicely with the Open311 API
	// FIXME: there HAS to be a better way to handle this.
	formatted_date_string := last_updated_at.Format(time.RFC3339)
	formatted_date_string_with_tz := formatted_date_string[0:len(formatted_date_string)-1] + "-0500" // trunc the trailing 'Z' and tack on timezone

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_after=" + formatted_date_string_with_tz

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests
}
<commit_msg>poll on a 30 second interval<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	_ "github.com/bmizerany/pq"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

const OPEN311_API_URI = "http://311api.cityofchicago.org/open311/v2/requests.json?extensions=true&page_size=500"

type Open311Request struct {
	Lat, Long                                                                                float64
	Ward, Police_district                                                                                int
	Service_request_id, Status, Service_name, Service_code, Agency_responsible, Address, Channel, Media_url string
	Requested_datetime, Updated_datetime                                                                    string // FIXME: should these be proper time objects?
	Extended_attributes                                                                                map[string]interface{}
}

func main() {
	// open database
	db, err := sql.Open("postgres", "dbname=cwfy sslmode=disable")
	if err != nil {
		log.Fatal("Cannot open database connection", err)
	}
	defer db.Close()

	last_run_at := time.Now()

	for {
		switch {
		case time.Since(last_run_at) > (30 * time.Second):
			poll_open311(db)
			last_run_at = time.Now()
		default:
			log.Print("sleeping for 10 seconds")
			time.Sleep(10 * time.Second)
		}
	}
}

func (req Open311Request) String() string {
	// pretty print SR information
	return fmt.Sprintf("%s: %s at %s %f,%f, last update %s", req.Service_request_id, req.Service_name, req.Address, req.Lat, req.Long, req.Updated_datetime)
}

func fetchRequests() (requests []Open311Request) {
	db, err := sql.Open("postgres", "dbname=cwfy sslmode=disable")
	if err != nil {
		log.Fatal("Cannot open database connection", err)
	}
	defer db.Close()

	// find the most recent SR that we know about in the database
	rows, err := db.Query("SELECT MAX(updated_datetime) FROM service_requests;")
	if err != nil {
		log.Fatal("error finding most recent service request", err)
	}

	last_updated_at := time.Now()
	for rows.Next() {
		if err := rows.Scan(&last_updated_at); err != nil {
			log.Fatal(err)
		}

		log.Printf("most recent SR timestamp %s", last_updated_at)
	}

	// janky hack to transform the last updated timestamp into
	// a format that plays nicely with the Open311 API
	// FIXME: there HAS to be a better way to handle this.
	formatted_date_string := last_updated_at.Format(time.RFC3339)
	formatted_date_string_with_tz := formatted_date_string[0:len(formatted_date_string)-1] + "-0500" // trunc the trailing 'Z' and tack on timezone

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_after=" + formatted_date_string_with_tz

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests
}

func poll_open311(db *sql.DB) {
	// load requests from open311
	requests := fetchRequests()

	for _, request := range requests {
		// for each request, either create or update the
		// corresponding record in the database.

		if request.Service_request_id == "" {
			log.Printf("Ignoring a request type %s because there is no SR number assigned", request.Service_name)
			continue
		}

		insert_stmt, err := db.Prepare("INSERT INTO service_requests(service_request_id," +
			"status, service_name, service_code, agency_responsible, " +
			"address, requested_datetime, updated_datetime, lat, long," +
			"ward, police_district, media_url, channel, duplicate, parent_service_request_id) " +
			"SELECT $1::varchar, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16 " +
			"WHERE NOT EXISTS (SELECT 1 FROM service_requests WHERE service_request_id = $1);")

		if err != nil {
			log.Fatal("error preparing database insert statement", err)
		}

		update_stmt, err := db.Prepare("UPDATE service_requests SET " +
			"status = $2, service_name = $3, service_code = $4, agency_responsible = $5, " +
			"address = $6, requested_datetime = $7, updated_datetime = $8, lat = $9, long = $10," +
			"ward = $11, police_district = $12, media_url = $13, channel = $14, duplicate = $15, " +
			"parent_service_request_id = $16, updated_at = NOW() WHERE service_request_id = $1;")

		if err != nil {
			log.Fatal("error preparing database update statement", err)
		}

		tx, err := db.Begin()

		if err != nil {
			log.Fatal("error beginning transaction", err)
		}

		_, err = tx.Stmt(update_stmt).Exec(request.Service_request_id,
			request.Status,
			request.Service_name,
			request.Service_code,
			request.Agency_responsible,
			request.Address,
			request.Requested_datetime,
			request.Updated_datetime,
			request.Lat,
			request.Long,
			request.Extended_attributes["ward"],
			request.Extended_attributes["police_district"],
			request.Media_url,
			request.Extended_attributes["channel"],
			request.Extended_attributes["duplicate"],
			request.Extended_attributes["parent_service_request_id"])

		if err != nil {
			log.Fatalf("could not update %s because %s", request.Service_request_id, err)
		}

		_, err = tx.Stmt(insert_stmt).Exec(request.Service_request_id,
			request.Status,
			request.Service_name,
			request.Service_code,
			request.Agency_responsible,
			request.Address,
			request.Requested_datetime,
			request.Updated_datetime,
			request.Lat,
			request.Long,
			request.Extended_attributes["ward"],
			request.Extended_attributes["police_district"],
			request.Media_url,
			request.Extended_attributes["channel"],
			request.Extended_attributes["duplicate"],
			request.Extended_attributes["parent_service_request_id"])

		if err != nil {
			log.Fatalf("could not save %s because %s", request.Service_request_id, err)
		} else {
			log.Printf("saved SR %s", request)
		}

		err = tx.Commit()
		if err != nil {
			log.Fatal("error closing transaction", err)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
)

const OPEN311_API_URI = "http://311api.cityofchicago.org/open311/v2/requests.json"

type Open311Request struct {
	Lat, Long                                                                           float64
	Service_request_id, Status, Service_name, Service_code, Agency_responsible, Address string
	Requested_datetime, Updated_datetime                                                string // FIXME: should these be proper time objects?
}

func main() {
	fetchRequests()
}

func (req Open311Request) String() string {
	return fmt.Sprintf("%s: %s at %s %f,%f", req.Service_request_id, req.Service_name, req.Address, req.Lat, req.Long)
}

func fetchRequests() {
	log.Printf("fetching from %s", OPEN311_API_URI)
	resp, err := http.Get(OPEN311_API_URI)
	defer resp.Body.Close()

	if err == nil {
		log.Println("fetch succesful, reading response")
		body, err := ioutil.ReadAll(resp.Body)

		if err == nil {
			log.Println("loaded response body.")
			// parse into JSON

			var requests []Open311Request
			err := json.Unmarshal(body, &requests)
			if err != nil {
				log.Fatal("error parsing JSON:", err)
			}

			log.Printf("received %d requests from Open311", len(requests))

			for i, req := range requests {
				fmt.Println(i, req)
			}
		}
	} else {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}
}
<commit_msg>working write-to-database<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	_ "github.com/bmizerany/pq"
	"io/ioutil"
	"log"
	"net/http"
)

const OPEN311_API_URI = "http://311api.cityofchicago.org/open311/v2/requests.json"

type Open311Request struct {
	Lat, Long                                                                           float64
	Service_request_id, Status, Service_name, Service_code, Agency_responsible, Address string
	Requested_datetime, Updated_datetime                                                string // FIXME: should these be proper time objects?
}

func main() {
	db, err := sql.Open("postgres", "dbname=cwfy sslmode=disable")
	if err == nil {
		requests := fetchRequests()
		for _, request := range requests {
			stmt, err := db.Prepare("INSERT INTO service_requests(service_request_id, status, service_name, service_code, agency_responsible, address, requested_datetime, updated_datetime, lat, long) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10);")
			if err != nil {
				log.Fatal("error preparing database statement", err)
			}

			_, err = stmt.Exec(request.Service_request_id,
				request.Status,
				request.Service_name,
				request.Service_code,
				request.Agency_responsible,
				request.Address,
				request.Requested_datetime,
				request.Updated_datetime,
				request.Lat,
				request.Long)

			if err != nil {
				log.Fatalf("could not save %s because %s", request.Service_request_id, err)
			}
		}
	} else {
		log.Fatal("Cannot open database connection", err)
	}
}

func (req Open311Request) String() string {
	return fmt.Sprintf("%s: %s at %s %f,%f", req.Service_request_id, req.Service_name, req.Address, req.Lat, req.Long)
}

// func (req Open311Request) Save() (saved bool) {
// 	stmt, err := db.Prepare("INSERT INTO service_requests (service_request_id, status, service_name, service_code, agency_responsible, address, requested_datetimeupdated_datetime, lat, long) VALUES (?,?,?,?,?,?,?);")
// 	if err != nil {
// 		log.Fatal("error preparing database statement", err)
// 	}
//
// 	_, err = stmt.Exec(req.Service_request_id,
// 		req.Status,
// 		req.Service_name,
// 		req.Service_code,
// 		req.Agency_responsible,
// 		req.Address,
// 		req.Requested_datetime,
// 		req.Updated_datetime,
// 		req.Lat,
// 		req.Long)
//
// 	if err != nil {
// 		log.Fatalf("could not save %s because %s", req.Service_request_id, err)
// 	}
//
// 	return true
// }

func fetchRequests() (requests []Open311Request) {
	log.Printf("fetching from %s", OPEN311_API_URI)
	resp, err := http.Get(OPEN311_API_URI)
	defer resp.Body.Close()

	if err == nil {
		log.Println("fetch succesful, reading response")
		body, err := ioutil.ReadAll(resp.Body)

		if err == nil {
			log.Println("loaded response body.")
			// parse into JSON

			// var requests []Open311Request
			err := json.Unmarshal(body, &requests)
			if err != nil {
				log.Fatal("error parsing JSON:", err)
			}

			log.Printf("received %d requests from Open311", len(requests))

		}
	} else {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	return requests
}
<|endoftext|>


<commit_before>package main_test

import (
	"time"

	. "github.com/alphagov/govuk_crawler_worker"
	. "github.com/alphagov/govuk_crawler_worker/queue"
	. "github.com/alphagov/govuk_crawler_worker/ttl_hash_set"

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

	"github.com/fzzy/radix/redis"
)

var _ = Describe("Workflow", func() {
	Describe("Acknowledging items", func() {
		exchangeName, queueName := "test-workflow-exchange", "test-workflow-queue"
		prefix := "govuk_mirror_crawler_workflow_test"

		var (
			queueManager    *QueueManager
			queueManagerErr error
			ttlHashSet      *TTLHashSet
			ttlHashSetErr   error
		)

		BeforeEach(func() {
			ttlHashSet, ttlHashSetErr = NewTTLHashSet(prefix, "127.0.0.1:6379")
			Expect(ttlHashSetErr).To(BeNil())

			queueManager, queueManagerErr = NewQueueManager(
				"amqp://guest:guest@localhost:5672/",
				exchangeName,
				queueName)

			Expect(queueManagerErr).To(BeNil())
			Expect(queueManager).ToNot(BeNil())
		})

		AfterEach(func() {
			Expect(ttlHashSet.Close()).To(BeNil())
			Expect(purgeAllKeys(prefix, "127.0.0.1:6379")).To(BeNil())

			deleted, err := queueManager.Consumer.Channel.QueueDelete(queueName, false, false, true)
			Expect(err).To(BeNil())
			Expect(deleted).To(Equal(0))

			err = queueManager.Consumer.Channel.ExchangeDelete(exchangeName, false, true)
			Expect(err).To(BeNil())

			queueManager.Close()
		})

		It("should read from a channel and add URLs to the hash set", func() {
			url := "https://www.gov.uk/foo"

			exists, err := ttlHashSet.Exists(url)
			Expect(err).To(BeNil())
			Expect(exists).To(BeFalse())

			deliveries, err := queueManager.Consume()
			Expect(err).To(BeNil())

			outbound := make(chan *CrawlerMessageItem, 1)

			err = queueManager.Publish("#", "text/plain", url)
			Expect(err).To(BeNil())

			for item := range deliveries {
				outbound <- NewCrawlerMessageItem(item, "www.gov.uk", []string{})
				break
			}

			Expect(len(outbound)).To(Equal(1))

			go AcknowledgeItem(outbound, ttlHashSet)
			time.Sleep(time.Millisecond)

			Expect(len(outbound)).To(Equal(0))

			exists, err = ttlHashSet.Exists(url)
			Expect(err).To(BeNil())
			Expect(exists).To(BeTrue())

			// Close the channel to stop the goroutine for AcknowledgeItem.
			close(outbound)
		})
	})
})

func purgeAllKeys(prefix string, address string) error {
	client, err := redis.Dial("tcp", address)
	if err != nil {
		return err
	}

	keys, err := client.Cmd("KEYS", prefix+"*").List()
	if err != nil || len(keys) <= 0 {
		return err
	}

	reply := client.Cmd("DEL", keys)
	if reply.Err != nil {
		return reply.Err
	}

	return nil
}
<commit_msg>Bound test block for AcknowledgeItem<commit_after>package main_test

import (
	"time"

	. "github.com/alphagov/govuk_crawler_worker"
	. "github.com/alphagov/govuk_crawler_worker/queue"
	. "github.com/alphagov/govuk_crawler_worker/ttl_hash_set"

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

	"github.com/fzzy/radix/redis"
)

var _ = Describe("Workflow", func() {
	Describe("Acknowledging items", func() {
		exchangeName, queueName := "test-workflow-exchange", "test-workflow-queue"
		prefix := "govuk_mirror_crawler_workflow_test"

		var (
			queueManager    *QueueManager
			queueManagerErr error
			ttlHashSet      *TTLHashSet
			ttlHashSetErr   error
		)

		BeforeEach(func() {
			ttlHashSet, ttlHashSetErr = NewTTLHashSet(prefix, "127.0.0.1:6379")
			Expect(ttlHashSetErr).To(BeNil())

			queueManager, queueManagerErr = NewQueueManager(
				"amqp://guest:guest@localhost:5672/",
				exchangeName,
				queueName)

			Expect(queueManagerErr).To(BeNil())
			Expect(queueManager).ToNot(BeNil())
		})

		AfterEach(func() {
			Expect(ttlHashSet.Close()).To(BeNil())
			Expect(purgeAllKeys(prefix, "127.0.0.1:6379")).To(BeNil())

			deleted, err := queueManager.Consumer.Channel.QueueDelete(queueName, false, false, true)
			Expect(err).To(BeNil())
			Expect(deleted).To(Equal(0))

			err = queueManager.Consumer.Channel.ExchangeDelete(exchangeName, false, true)
			Expect(err).To(BeNil())

			queueManager.Close()
		})

		Describe("AcknowledgeItem", func() {
			It("should read from a channel and add URLs to the hash set", func() {
				url := "https://www.gov.uk/foo"

				exists, err := ttlHashSet.Exists(url)
				Expect(err).To(BeNil())
				Expect(exists).To(BeFalse())

				deliveries, err := queueManager.Consume()
				Expect(err).To(BeNil())

				outbound := make(chan *CrawlerMessageItem, 1)

				err = queueManager.Publish("#", "text/plain", url)
				Expect(err).To(BeNil())

				for item := range deliveries {
					outbound <- NewCrawlerMessageItem(item, "www.gov.uk", []string{})
					break
				}

				Expect(len(outbound)).To(Equal(1))

				go AcknowledgeItem(outbound, ttlHashSet)
				time.Sleep(time.Millisecond)

				Expect(len(outbound)).To(Equal(0))

				exists, err = ttlHashSet.Exists(url)
				Expect(err).To(BeNil())
				Expect(exists).To(BeTrue())

				// Close the channel to stop the goroutine for AcknowledgeItem.
				close(outbound)
			})
		})
	})
})

func purgeAllKeys(prefix string, address string) error {
	client, err := redis.Dial("tcp", address)
	if err != nil {
		return err
	}

	keys, err := client.Cmd("KEYS", prefix+"*").List()
	if err != nil || len(keys) <= 0 {
		return err
	}

	reply := client.Cmd("DEL", keys)
	if reply.Err != nil {
		return reply.Err
	}

	return nil
}
<|endoftext|>


<commit_before>package nyb

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/url"
	"strings"
	"time"

	"github.com/hako/durafmt"
	"github.com/ugjka/dumbirc"
	"github.com/ugjka/go-tz"
)

var nickChangeInterval = time.Second * 5

func (bot *Settings) addCallbacks() {
	irc := bot.IrcConn
	//On any message send a signal to ping timer to be ready

	irc.AddCallback(dumbirc.WELCOME, func(msg *dumbirc.Message) {
		if irc.Password != "" {
			confirmErr := fmt.Errorf("did not get identification confirmation")
			err := irc.WaitFor(func(m *dumbirc.Message) bool {
				return m.Command == dumbirc.NOTICE && strings.Contains(m.Content, "You are now identified for")
			},
				func() {},
				time.Second*30,
				confirmErr,
			)
			if err == confirmErr {
				log.Println(err)
				log.Println("trying to start the bot anyway")
			} else if err != nil {
				return
			}
		}
		irc.Join(bot.IrcChans)
		//Prevent early start
		bot.Do(func() {
			close(bot.start)
		})
	})
}

func (bot *Settings) addTriggers() {
	irc := bot.IrcConn
	//Trigger for !help
	stHelp := "%s: Query location: '%s <location>', Next zone: '%s !next', Last zone: '%s !last', Remaining: '%s !remaining', Source code: https://github.com/ugjka/newyearsbot"
	irc.AddTrigger(dumbirc.Trigger{
		Condition: func(msg *dumbirc.Message) bool {
			return msg.Command == dumbirc.PRIVMSG &&
				msg.Content == fmt.Sprintf("%s !help", bot.IrcTrigger)
		},
		Response: func(msg *dumbirc.Message) {
			log.Println("Querying !help...")
			irc.Reply(msg, fmt.Sprintf(stHelp, msg.Name, bot.IrcTrigger, bot.IrcTrigger, bot.IrcTrigger, bot.IrcTrigger))
		},
	})
	//Trigger for !next
	irc.AddTrigger(dumbirc.Trigger{
		Condition: func(msg *dumbirc.Message) bool {
			return msg.Command == dumbirc.PRIVMSG &&
				msg.Content == fmt.Sprintf("%s !next", bot.IrcTrigger)
		},
		Response: func(msg *dumbirc.Message) {
			log.Println("Querying !next...")
			dur := time.Minute * time.Duration(bot.next.Offset*60)
			if timeNow().UTC().Add(dur).After(target) {
				irc.Reply(msg, fmt.Sprintf("No more next, %d is here AoE", target.Year()))
				return
			}
			humandur := durafmt.Parse(target.Sub(timeNow().UTC().Add(dur)))
			irc.Reply(msg, fmt.Sprintf("Next New Year in %s in %s",
				removeMilliseconds(humandur), bot.next))
		},
	})
	//Trigger for !last
	irc.AddTrigger(dumbirc.Trigger{
		Condition: func(msg *dumbirc.Message) bool {
			return msg.Command == dumbirc.PRIVMSG &&
				msg.Content == fmt.Sprintf("%s !last", bot.IrcTrigger)
		},
		Response: func(msg *dumbirc.Message) {
			log.Println("Querying !last...")
			dur := time.Minute * time.Duration(bot.last.Offset*60)
			humandur := durafmt.Parse(timeNow().UTC().Add(dur).Sub(target))
			if bot.last.Offset == -12 {
				humandur = durafmt.Parse(timeNow().UTC().Add(dur).Sub(target.AddDate(-1, 0, 0)))
			}
			irc.Reply(msg, fmt.Sprintf("Last New Year %s ago in %s",
				removeMilliseconds(humandur), bot.last))
		},
	})
	//Trigger for !remaining
	irc.AddTrigger(dumbirc.Trigger{
		Condition: func(msg *dumbirc.Message) bool {
			return msg.Command == dumbirc.PRIVMSG &&
				msg.Content == fmt.Sprintf("%s !remaining", bot.IrcTrigger)
		},
		Response: func(msg *dumbirc.Message) {
			log.Println("Querying !remaining...")
			ss := "s"
			if bot.remaining == 1 {
				ss = ""
			}
			irc.Reply(msg, fmt.Sprintf("%s: %d timezone%s remaining", msg.Name, bot.remaining, ss))
		},
	})
	//Trigger for location queries
	irc.AddTrigger(dumbirc.Trigger{
		Condition: func(msg *dumbirc.Message) bool {
			return msg.Command == dumbirc.PRIVMSG &&
				!strings.Contains(msg.Content, "!next") &&
				!strings.Contains(msg.Content, "!last") &&
				!strings.Contains(msg.Content, "!help") &&
				!strings.Contains(msg.Content, "!remaining") &&
				strings.HasPrefix(msg.Content, fmt.Sprintf("%s ", bot.IrcTrigger))
		},
		Response: func(msg *dumbirc.Message) {
			tz, err := bot.getNewYear(msg.Content[len(bot.IrcTrigger)+1:])
			if err == errNoZone || err == errNoPlace {
				log.Println("Query error:", err)
				irc.Reply(msg, fmt.Sprintf("%s: %s", msg.Name, err))
				return
			}
			if err != nil {
				log.Println("Query error:", err)
				irc.Reply(msg, fmt.Sprintf("%s: Some error occurred!", msg.Name))
				return
			}
			irc.Reply(msg, fmt.Sprintf("%s: %s", msg.Name, tz))
		},
	})
}

var (
	errNoZone  = errors.New("couldn't get the timezone for that location")
	errNoPlace = errors.New("Couldn't find that place")
)

func (bot *Settings) getNominatimReqURL(location *string) string {
	maps := url.Values{}
	maps.Add("q", *location)
	maps.Add("format", "json")
	maps.Add("accept-language", "en")
	maps.Add("limit", "1")
	maps.Add("email", bot.Email)
	return bot.Nominatim + NominatimEndpoint + maps.Encode()
}

var stNewYearWillHappen = "New Year in %s will happen in %s"
var stNewYearHappenned = "New Year in %s happened %s ago"

func (bot *Settings) getNewYear(location string) (string, error) {
	log.Println("Querying location:", location)
	data, err := NominatimGetter(bot.getNominatimReqURL(&location))
	if err != nil {
		log.Println(err)
		return "", err
	}
	var res NominatimResults
	if err = json.Unmarshal(data, &res); err != nil {
		log.Println(err)
		return "", err
	}
	if len(res) == 0 {
		return "", errNoPlace
	}
	p := gotz.Point{
		Lat: res[0].Lat,
		Lon: res[0].Lon,
	}
	tzid, err := gotz.GetZone(p)
	if err != nil {
		return "", errNoZone
	}
	zone, err := time.LoadLocation(tzid[0])
	if err != nil {
		return "", errNoZone
	}
	offset := time.Second * time.Duration(getOffset(target, zone))
	address := res[0].DisplayName

	if timeNow().UTC().Add(offset).Before(target) {
		humandur := durafmt.Parse(target.Sub(timeNow().UTC().Add(offset)))
		return fmt.Sprintf(stNewYearWillHappen, address, removeMilliseconds(humandur)), nil
	}
	humandur := durafmt.Parse(timeNow().UTC().Add(offset).Sub(target))
	return fmt.Sprintf(stNewYearHappenned, address, removeMilliseconds(humandur)), nil
}
<commit_msg>add time in location<commit_after>package nyb

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/url"
	"strings"
	"time"

	"github.com/hako/durafmt"
	"github.com/ugjka/dumbirc"
	"github.com/ugjka/go-tz"
)

var nickChangeInterval = time.Second * 5

func (bot *Settings) addCallbacks() {
	irc := bot.IrcConn
	//On any message send a signal to ping timer to be ready

	irc.AddCallback(dumbirc.WELCOME, func(msg *dumbirc.Message) {
		if irc.Password != "" {
			confirmErr := fmt.Errorf("did not get identification confirmation")
			err := irc.WaitFor(func(m *dumbirc.Message) bool {
				return m.Command == dumbirc.NOTICE && strings.Contains(m.Content, "You are now identified for")
			},
				func() {},
				time.Second*30,
				confirmErr,
			)
			if err == confirmErr {
				log.Println(err)
				log.Println("trying to start the bot anyway")
			} else if err != nil {
				return
			}
		}
		irc.Join(bot.IrcChans)
		//Prevent early start
		bot.Do(func() {
			close(bot.start)
		})
	})
}

func (bot *Settings) addTriggers() {
	irc := bot.IrcConn
	//Trigger for !help
	stHelp := "%s: Query location: '%s <location>', Time in location: '%s !time <location>', Next zone: '%s !next', Last zone: '%s !last', Remaining: '%s !remaining', Source code: https://github.com/ugjka/newyearsbot"
	irc.AddTrigger(dumbirc.Trigger{
		Condition: func(msg *dumbirc.Message) bool {
			return msg.Command == dumbirc.PRIVMSG &&
				msg.Content == fmt.Sprintf("%s !help", bot.IrcTrigger)
		},
		Response: func(msg *dumbirc.Message) {
			log.Println("Querying !help...")
			irc.Reply(msg, fmt.Sprintf(stHelp, msg.Name, bot.IrcTrigger, bot.IrcTrigger, bot.IrcTrigger, bot.IrcTrigger, bot.IrcTrigger))
		},
	})
	//Trigger for !next
	irc.AddTrigger(dumbirc.Trigger{
		Condition: func(msg *dumbirc.Message) bool {
			return msg.Command == dumbirc.PRIVMSG &&
				msg.Content == fmt.Sprintf("%s !next", bot.IrcTrigger)
		},
		Response: func(msg *dumbirc.Message) {
			log.Println("Querying !next...")
			dur := time.Minute * time.Duration(bot.next.Offset*60)
			if timeNow().UTC().Add(dur).After(target) {
				irc.Reply(msg, fmt.Sprintf("No more next, %d is here AoE", target.Year()))
				return
			}
			humandur := durafmt.Parse(target.Sub(timeNow().UTC().Add(dur)))
			irc.Reply(msg, fmt.Sprintf("Next New Year in %s in %s",
				removeMilliseconds(humandur), bot.next))
		},
	})
	//Trigger for !last
	irc.AddTrigger(dumbirc.Trigger{
		Condition: func(msg *dumbirc.Message) bool {
			return msg.Command == dumbirc.PRIVMSG &&
				msg.Content == fmt.Sprintf("%s !last", bot.IrcTrigger)
		},
		Response: func(msg *dumbirc.Message) {
			log.Println("Querying !last...")
			dur := time.Minute * time.Duration(bot.last.Offset*60)
			humandur := durafmt.Parse(timeNow().UTC().Add(dur).Sub(target))
			if bot.last.Offset == -12 {
				humandur = durafmt.Parse(timeNow().UTC().Add(dur).Sub(target.AddDate(-1, 0, 0)))
			}
			irc.Reply(msg, fmt.Sprintf("Last New Year %s ago in %s",
				removeMilliseconds(humandur), bot.last))
		},
	})
	//Trigger for !remaining
	irc.AddTrigger(dumbirc.Trigger{
		Condition: func(msg *dumbirc.Message) bool {
			return msg.Command == dumbirc.PRIVMSG &&
				msg.Content == fmt.Sprintf("%s !remaining", bot.IrcTrigger)
		},
		Response: func(msg *dumbirc.Message) {
			log.Println("Querying !remaining...")
			ss := "s"
			if bot.remaining == 1 {
				ss = ""
			}
			irc.Reply(msg, fmt.Sprintf("%s: %d timezone%s remaining", msg.Name, bot.remaining, ss))
		},
	})
	//Trigger for time in location
	irc.AddTrigger(dumbirc.Trigger{
		Condition: func(msg *dumbirc.Message) bool {
			return msg.Command == dumbirc.PRIVMSG &&
				strings.HasPrefix(msg.Content, fmt.Sprintf("%s !time ", bot.IrcTrigger))
		},
		Response: func(msg *dumbirc.Message) {
			log.Println("Querying !time...")
			res, err := bot.getTime(msg.Content[9:])
			if err == errNoZone || err == errNoPlace {
				log.Println("Query error:", err)
				irc.Reply(msg, fmt.Sprintf("%s: %s", msg.Name, err))
				return
			}
			if err != nil {
				log.Println("Query error:", err)
				irc.Reply(msg, fmt.Sprintf("%s: Some error occurred!", msg.Name))
				return
			}
			irc.Reply(msg, res)
		},
	})
	//Trigger for location queries
	irc.AddTrigger(dumbirc.Trigger{
		Condition: func(msg *dumbirc.Message) bool {
			return msg.Command == dumbirc.PRIVMSG &&
				!strings.Contains(msg.Content, "!time") &&
				!strings.Contains(msg.Content, "!next") &&
				!strings.Contains(msg.Content, "!last") &&
				!strings.Contains(msg.Content, "!help") &&
				!strings.Contains(msg.Content, "!remaining") &&
				strings.HasPrefix(msg.Content, fmt.Sprintf("%s ", bot.IrcTrigger))
		},
		Response: func(msg *dumbirc.Message) {
			tz, err := bot.getNewYear(msg.Content[len(bot.IrcTrigger)+1:])
			if err == errNoZone || err == errNoPlace {
				log.Println("Query error:", err)
				irc.Reply(msg, fmt.Sprintf("%s: %s", msg.Name, err))
				return
			}
			if err != nil {
				log.Println("Query error:", err)
				irc.Reply(msg, fmt.Sprintf("%s: Some error occurred!", msg.Name))
				return
			}
			irc.Reply(msg, fmt.Sprintf("%s: %s", msg.Name, tz))
		},
	})
}

var (
	errNoZone  = errors.New("couldn't get the timezone for that location")
	errNoPlace = errors.New("Couldn't find that place")
)

func (bot *Settings) getNominatimReqURL(location *string) string {
	maps := url.Values{}
	maps.Add("q", *location)
	maps.Add("format", "json")
	maps.Add("accept-language", "en")
	maps.Add("limit", "1")
	maps.Add("email", bot.Email)
	return bot.Nominatim + NominatimEndpoint + maps.Encode()
}

var stNewYearWillHappen = "New Year in %s will happen in %s"
var stNewYearHappenned = "New Year in %s happened %s ago"

func (bot *Settings) getTime(location string) (string, error) {
	log.Println("Querying location:", location)
	data, err := NominatimGetter(bot.getNominatimReqURL(&location))
	if err != nil {
		log.Println(err)
		return "", err
	}
	var res NominatimResults
	if err = json.Unmarshal(data, &res); err != nil {
		log.Println(err)
		return "", err
	}
	if len(res) == 0 {
		return "", errNoPlace
	}
	p := gotz.Point{
		Lat: res[0].Lat,
		Lon: res[0].Lon,
	}
	tzid, err := gotz.GetZone(p)
	if err != nil {
		return "", errNoZone
	}
	zone, err := time.LoadLocation(tzid[0])
	if err != nil {
		return "", errNoZone
	}
	address := res[0].DisplayName
	return fmt.Sprintf("Time in %s is %s", address, time.Now().In(zone).Format("Mon Jan 2 15:04:05 -0700 MST 2006")), nil
}

func (bot *Settings) getNewYear(location string) (string, error) {
	log.Println("Querying location:", location)
	data, err := NominatimGetter(bot.getNominatimReqURL(&location))
	if err != nil {
		log.Println(err)
		return "", err
	}
	var res NominatimResults
	if err = json.Unmarshal(data, &res); err != nil {
		log.Println(err)
		return "", err
	}
	if len(res) == 0 {
		return "", errNoPlace
	}
	p := gotz.Point{
		Lat: res[0].Lat,
		Lon: res[0].Lon,
	}
	tzid, err := gotz.GetZone(p)
	if err != nil {
		return "", errNoZone
	}
	zone, err := time.LoadLocation(tzid[0])
	if err != nil {
		return "", errNoZone
	}
	offset := time.Second * time.Duration(getOffset(target, zone))
	address := res[0].DisplayName

	if timeNow().UTC().Add(offset).Before(target) {
		humandur := durafmt.Parse(target.Sub(timeNow().UTC().Add(offset)))
		return fmt.Sprintf(stNewYearWillHappen, address, removeMilliseconds(humandur)), nil
	}
	humandur := durafmt.Parse(timeNow().UTC().Add(offset).Sub(target))
	return fmt.Sprintf(stNewYearHappenned, address, removeMilliseconds(humandur)), nil
}
<|endoftext|>


<commit_before>// package container contains convenience functions for manipulating containers
package container

import (
	"errors"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/fsouza/go-dockerclient"
)

const defaultDockerHost = "unix:///var/run/docker.sock"

var DefaultConfig = docker.Config{
	AttachStdin:  true,
	AttachStdout: true,
	AttachStderr: true,
	OpenStdin:    true,
	StdinOnce:    true,
}

func RawStartContainer(opts docker.CreateContainerOptions) (string, error) {
	client, err := getNewDockerClient()
	if err != nil {
		log.Print(err)
		return "", err
	}
	container, err := client.CreateContainer(opts)
	if err != nil {
		log.Print(err)
		return "", err
	}
	err = client.StartContainer(container.ID, opts.HostConfig)
	if err != nil {
		log.Print(err)
		return "", err
	}

	return container.ID, nil
}

func StartContainer(image string, command []string) (string, error) {
	config := docker.Config{Image: image, Cmd: command}
	opts := docker.CreateContainerOptions{Config: &config}
	return RawStartContainer(opts)
}

func StopContainer(id string) error {
	client, err := getNewDockerClient()
	if err != nil {
		log.Print(err)
		return err
	}
	return client.StopContainer(id, 5)
}

func KillContainer(id string) error {
	client, err := getNewDockerClient()
	if err != nil {
		log.Print(err)
		return err
	}
	return client.KillContainer(docker.KillContainerOptions{ID: id})
}

func IpAddr(containerId string) (string, error) {
	client, err := getNewDockerClient()
	if err != nil {
		log.Print(err)
		return "", err
	}
	container, err := client.InspectContainer(containerId)
	if err != nil {
		log.Print(err)
		return "", err
	}

	return container.NetworkSettings.IPAddress, nil
}

func PullImage(image string) error {
	repo_tag := strings.Split(image, ":")
	client, err := getNewDockerClient()
	if err != nil {
		log.Print(err)
		return err
	}
	opts := docker.PullImageOptions{Repository: repo_tag[0], Tag: "latest"}
	if len(repo_tag) == 2 {
		opts.Tag = repo_tag[1]
	}
	return client.PullImage(opts, docker.AuthConfiguration{})
}

func PipeToStdin(id string, in io.Reader) error {
	client, err := getNewDockerClient()
	if err != nil {
		log.Print(err)
		return err
	}
	return client.AttachToContainer(docker.AttachToContainerOptions{
		Container:   id,
		InputStream: in,
		Stdin:       true,
		Stream:      true,
	})
}

func ContainerLogs(id string, out io.Writer) error {
	client, err := getNewDockerClient()
	if err != nil {
		log.Print(err)
		return err
	}
	return client.AttachToContainer(docker.AttachToContainerOptions{
		Container:    id,
		OutputStream: out,
		ErrorStream:  out,
		Stdout:       true,
		Stderr:       true,
		Logs:         true,
	})
}

func WaitContainer(id string) (int, error) {
	client, err := getNewDockerClient()
	if err != nil {
		log.Print(err)
		return 0, err
	}

	return client.WaitContainer(id)
}

func getNewDockerClient() (*docker.Client, error) {
	if os.Getenv("DOCKER_TLS_VERIFY") == "1" {
		path := os.Getenv("DOCKER_CERT_PATH")

		if path == "" {
			path = os.Getenv("HOME")

			if path == "" {
				return nil, errors.New("pfs: environment variable HOME must be set if DOCKER_CERT_PATH is not set")
			}

			var err error

			path = filepath.Join(path, ".docker")
			path, err = filepath.Abs(path)

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

		return docker.NewTLSClient(
			getDockerHost(),
			path+"/cert.pem",
			path+"/key.pem",
			path+"/ca.pem",
		)
	}

	return docker.NewClient(getDockerHost())
}

func getDockerHost() (host string) {
	host = os.Getenv("DOCKER_HOST")

	if host == "" {
		host = defaultDockerHost
	}

	return
}
<commit_msg>Rename the Docker client creation function and merge with the Docker host function<commit_after>// package container contains convenience functions for manipulating containers
package container

import (
	"errors"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/fsouza/go-dockerclient"
)

const defaultDockerHost = "unix:///var/run/docker.sock"

var DefaultConfig = docker.Config{
	AttachStdin:  true,
	AttachStdout: true,
	AttachStderr: true,
	OpenStdin:    true,
	StdinOnce:    true,
}

func RawStartContainer(opts docker.CreateContainerOptions) (string, error) {
	client, err := NewDockerClientFromEnv()
	if err != nil {
		log.Print(err)
		return "", err
	}
	container, err := client.CreateContainer(opts)
	if err != nil {
		log.Print(err)
		return "", err
	}
	err = client.StartContainer(container.ID, opts.HostConfig)
	if err != nil {
		log.Print(err)
		return "", err
	}

	return container.ID, nil
}

func StartContainer(image string, command []string) (string, error) {
	config := docker.Config{Image: image, Cmd: command}
	opts := docker.CreateContainerOptions{Config: &config}
	return RawStartContainer(opts)
}

func StopContainer(id string) error {
	client, err := NewDockerClientFromEnv()
	if err != nil {
		log.Print(err)
		return err
	}
	return client.StopContainer(id, 5)
}

func KillContainer(id string) error {
	client, err := NewDockerClientFromEnv()
	if err != nil {
		log.Print(err)
		return err
	}
	return client.KillContainer(docker.KillContainerOptions{ID: id})
}

func IpAddr(containerId string) (string, error) {
	client, err := NewDockerClientFromEnv()
	if err != nil {
		log.Print(err)
		return "", err
	}
	container, err := client.InspectContainer(containerId)
	if err != nil {
		log.Print(err)
		return "", err
	}

	return container.NetworkSettings.IPAddress, nil
}

func PullImage(image string) error {
	repo_tag := strings.Split(image, ":")
	client, err := NewDockerClientFromEnv()
	if err != nil {
		log.Print(err)
		return err
	}
	opts := docker.PullImageOptions{Repository: repo_tag[0], Tag: "latest"}
	if len(repo_tag) == 2 {
		opts.Tag = repo_tag[1]
	}
	return client.PullImage(opts, docker.AuthConfiguration{})
}

func PipeToStdin(id string, in io.Reader) error {
	client, err := NewDockerClientFromEnv()
	if err != nil {
		log.Print(err)
		return err
	}
	return client.AttachToContainer(docker.AttachToContainerOptions{
		Container:   id,
		InputStream: in,
		Stdin:       true,
		Stream:      true,
	})
}

func ContainerLogs(id string, out io.Writer) error {
	client, err := NewDockerClientFromEnv()
	if err != nil {
		log.Print(err)
		return err
	}
	return client.AttachToContainer(docker.AttachToContainerOptions{
		Container:    id,
		OutputStream: out,
		ErrorStream:  out,
		Stdout:       true,
		Stderr:       true,
		Logs:         true,
	})
}

func WaitContainer(id string) (int, error) {
	client, err := NewDockerClientFromEnv()
	if err != nil {
		log.Print(err)
		return 0, err
	}

	return client.WaitContainer(id)
}

func NewDockerClientFromEnv() (*docker.Client, error) {
	host := os.Getenv("DOCKER_HOST")

	if host == "" {
		host = defaultDockerHost
	}

	if os.Getenv("DOCKER_TLS_VERIFY") != "" {
		path := os.Getenv("DOCKER_CERT_PATH")

		if path == "" {
			path = os.Getenv("HOME")

			if path == "" {
				return nil, errors.New("pfs: environment variable HOME must be set if DOCKER_CERT_PATH is not set")
			}

			var err error

			path = filepath.Join(path, ".docker")
			path, err = filepath.Abs(path)

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

		return docker.NewTLSClient(
			host,
			path+"/cert.pem",
			path+"/key.pem",
			path+"/ca.pem",
		)
	}

	return docker.NewClient(host)
}
<|endoftext|>


<commit_before>package oauth2

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

	"github.com/hashicorp/go-retryablehttp"

	"github.com/ory/hydra/x"

	"github.com/ory/fosite"
	"github.com/ory/hydra/consent"
	"github.com/ory/hydra/driver/config"
	"github.com/ory/x/errorsx"
)

// AccessRequestHook is called when an access token is being refreshed.
type AccessRequestHook func(ctx context.Context, requester fosite.AccessRequester) error

// Requester is a token endpoint's request context.
//
// swagger:model oAuth2AccessRequest
type Requester struct {
	// ClientID is the identifier of the OAuth 2.0 client.
	ClientID string `json:"client_id"`
	// GrantedScopes is the list of scopes granted to the OAuth 2.0 client.
	GrantedScopes []string `json:"granted_scopes"`
	// GrantedAudience is the list of audiences granted to the OAuth 2.0 client.
	GrantedAudience []string `json:"granted_audience"`
	// GrantTypes is the requests grant types.
	GrantTypes []string `json:"grant_types"`
}

// RefreshTokenHookRequest is the request body sent to the refresh token hook.
//
// swagger:model refreshTokenHookRequest
type RefreshTokenHookRequest struct {
	// Subject is the identifier of the authenticated end-user.
	Subject string `json:"subject"`
	// Session is the request's session..
	Session *Session `json:"session"`
	// Requester is a token endpoint's request context.
	Requester Requester `json:"requester"`
	// ClientID is the identifier of the OAuth 2.0 client.
	ClientID string `json:"client_id"`
	// GrantedScopes is the list of scopes granted to the OAuth 2.0 client.
	GrantedScopes []string `json:"granted_scopes"`
	// GrantedAudience is the list of audiences granted to the OAuth 2.0 client.
	GrantedAudience []string `json:"granted_audience"`
}

// RefreshTokenHookResponse is the response body received from the refresh token hook.
//
// swagger:model refreshTokenHookResponse
type RefreshTokenHookResponse struct {
	// Session is the session data returned by the hook.
	Session consent.ConsentRequestSessionData `json:"session"`
}

// RefreshTokenHook is an AccessRequestHook called for `refresh_token` grant type.
func RefreshTokenHook(reg interface {
	config.Provider
	x.HTTPClientProvider
}) AccessRequestHook {
	return func(ctx context.Context, requester fosite.AccessRequester) error {
		hookURL := reg.Config().TokenRefreshHookURL(ctx)
		if hookURL == nil {
			return nil
		}

		if !requester.GetGrantTypes().ExactOne("refresh_token") {
			return nil
		}

		session, ok := requester.GetSession().(*Session)
		if !ok {
			return nil
		}

		requesterInfo := Requester{
			ClientID:        requester.GetClient().GetID(),
			GrantedScopes:   requester.GetGrantedScopes(),
			GrantedAudience: requester.GetGrantedAudience(),
			GrantTypes:      requester.GetGrantTypes(),
		}

		reqBody := RefreshTokenHookRequest{
			Session:         session,
			Requester:       requesterInfo,
			Subject:         session.GetSubject(),
			ClientID:        requester.GetClient().GetID(),
			GrantedScopes:   requester.GetGrantedScopes(),
			GrantedAudience: requester.GetGrantedAudience(),
		}
		reqBodyBytes, err := json.Marshal(&reqBody)
		if err != nil {
			return errorsx.WithStack(
				fosite.ErrServerError.
					WithWrap(err).
					WithDebug("refresh token hook: marshal request body"),
			)
		}

		req, err := retryablehttp.NewRequestWithContext(ctx, http.MethodPost, hookURL.String(), bytes.NewReader(reqBodyBytes))
		if err != nil {
			return errorsx.WithStack(
				fosite.ErrServerError.
					WithWrap(err).
					WithDebug("refresh token hook: new http request"),
			)
		}
		req.Header.Set("Content-Type", "application/json; charset=UTF-8")

		resp, err := reg.HTTPClient(ctx).Do(req)
		if err != nil {
			return errorsx.WithStack(
				fosite.ErrServerError.
					WithWrap(err).
					WithDebug("refresh token hook: do http request"),
			)
		}
		defer resp.Body.Close()

		switch resp.StatusCode {
		case http.StatusOK:
			// Token refresh permitted with new session data
		case http.StatusNoContent:
			// Token refresh is permitted without overriding session data
			return nil
		case http.StatusForbidden:
			return errorsx.WithStack(
				fosite.ErrAccessDenied.
					WithDebugf("refresh token hook: %s", resp.Status),
			)
		default:
			return errorsx.WithStack(
				fosite.ErrServerError.
					WithDebugf("refresh token hook: %s", resp.Status),
			)
		}

		var respBody RefreshTokenHookResponse
		if err := json.NewDecoder(resp.Body).Decode(&respBody); err != nil {
			return errorsx.WithStack(
				fosite.ErrServerError.
					WithWrap(err).
					WithDebugf("refresh token hook: unmarshal response body"),
			)
		}

		// Overwrite existing session data (extra claims).
		session.Extra = respBody.Session.AccessToken
		idTokenClaims := session.IDTokenClaims()
		idTokenClaims.Extra = respBody.Session.IDToken

		return nil
	}
}
<commit_msg>feat: improve refresh token error messages<commit_after>package oauth2

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

	"github.com/hashicorp/go-retryablehttp"

	"github.com/ory/hydra/x"

	"github.com/ory/fosite"
	"github.com/ory/hydra/consent"
	"github.com/ory/hydra/driver/config"
	"github.com/ory/x/errorsx"
)

// AccessRequestHook is called when an access token is being refreshed.
type AccessRequestHook func(ctx context.Context, requester fosite.AccessRequester) error

// Requester is a token endpoint's request context.
//
// swagger:model oAuth2AccessRequest
type Requester struct {
	// ClientID is the identifier of the OAuth 2.0 client.
	ClientID string `json:"client_id"`
	// GrantedScopes is the list of scopes granted to the OAuth 2.0 client.
	GrantedScopes []string `json:"granted_scopes"`
	// GrantedAudience is the list of audiences granted to the OAuth 2.0 client.
	GrantedAudience []string `json:"granted_audience"`
	// GrantTypes is the requests grant types.
	GrantTypes []string `json:"grant_types"`
}

// RefreshTokenHookRequest is the request body sent to the refresh token hook.
//
// swagger:model refreshTokenHookRequest
type RefreshTokenHookRequest struct {
	// Subject is the identifier of the authenticated end-user.
	Subject string `json:"subject"`
	// Session is the request's session..
	Session *Session `json:"session"`
	// Requester is a token endpoint's request context.
	Requester Requester `json:"requester"`
	// ClientID is the identifier of the OAuth 2.0 client.
	ClientID string `json:"client_id"`
	// GrantedScopes is the list of scopes granted to the OAuth 2.0 client.
	GrantedScopes []string `json:"granted_scopes"`
	// GrantedAudience is the list of audiences granted to the OAuth 2.0 client.
	GrantedAudience []string `json:"granted_audience"`
}

// RefreshTokenHookResponse is the response body received from the refresh token hook.
//
// swagger:model refreshTokenHookResponse
type RefreshTokenHookResponse struct {
	// Session is the session data returned by the hook.
	Session consent.AcceptOAuth2ConsentRequestSession `json:"session"`
}

// RefreshTokenHook is an AccessRequestHook called for `refresh_token` grant type.
func RefreshTokenHook(reg interface {
	config.Provider
	x.HTTPClientProvider
}) AccessRequestHook {
	return func(ctx context.Context, requester fosite.AccessRequester) error {
		hookURL := reg.Config().TokenRefreshHookURL(ctx)
		if hookURL == nil {
			return nil
		}

		if !requester.GetGrantTypes().ExactOne("refresh_token") {
			return nil
		}

		session, ok := requester.GetSession().(*Session)
		if !ok {
			return nil
		}

		requesterInfo := Requester{
			ClientID:        requester.GetClient().GetID(),
			GrantedScopes:   requester.GetGrantedScopes(),
			GrantedAudience: requester.GetGrantedAudience(),
			GrantTypes:      requester.GetGrantTypes(),
		}

		reqBody := RefreshTokenHookRequest{
			Session:         session,
			Requester:       requesterInfo,
			Subject:         session.GetSubject(),
			ClientID:        requester.GetClient().GetID(),
			GrantedScopes:   requester.GetGrantedScopes(),
			GrantedAudience: requester.GetGrantedAudience(),
		}
		reqBodyBytes, err := json.Marshal(&reqBody)
		if err != nil {
			return errorsx.WithStack(
				fosite.ErrServerError.
					WithWrap(err).
					WithDescription("An error occurred while encoding the refresh token hook.").
					WithDebugf("Unable to encode the refresh token hook body: %s", err),
			)
		}

		req, err := retryablehttp.NewRequestWithContext(ctx, http.MethodPost, hookURL.String(), bytes.NewReader(reqBodyBytes))
		if err != nil {
			return errorsx.WithStack(
				fosite.ErrServerError.
					WithWrap(err).
					WithDescription("An error occurred while preparing the refresh token hook.").
					WithDebugf("Unable to prepare the HTTP Request: %s", err),
			)
		}
		req.Header.Set("Content-Type", "application/json; charset=UTF-8")

		resp, err := reg.HTTPClient(ctx).Do(req)
		if err != nil {
			return errorsx.WithStack(
				fosite.ErrServerError.
					WithWrap(err).
					WithDescription("An error occurred while executing the refresh token hook.").
					WithDebugf("Unable to execute HTTP Request: %s", err),
			)
		}
		defer resp.Body.Close()

		switch resp.StatusCode {
		case http.StatusOK:
			// Token refresh permitted with new session data
		case http.StatusNoContent:
			// Token refresh is permitted without overriding session data
			return nil
		case http.StatusForbidden:
			return errorsx.WithStack(
				fosite.ErrAccessDenied.
					WithDescription("The refresh token hook target responded with an error.").
					WithDebugf("Refresh token hook responded with HTTP status code: %s", resp.Status),
			)
		default:
			return errorsx.WithStack(
				fosite.ErrServerError.
					WithDescription("The refresh token hook target responded with an error.").
					WithDebugf("Refresh token hook responded with HTTP status code: %s", resp.Status),
			)
		}

		var respBody RefreshTokenHookResponse
		if err := json.NewDecoder(resp.Body).Decode(&respBody); err != nil {
			return errorsx.WithStack(
				fosite.ErrServerError.
					WithWrap(err).
					WithDescription("The refresh token hook target responded with an error.").
					WithDebugf("Response from refresh token hook could not be decoded: %s", err),
			)
		}

		// Overwrite existing session data (extra claims).
		session.Extra = respBody.Session.AccessToken
		idTokenClaims := session.IDTokenClaims()
		idTokenClaims.Extra = respBody.Session.IDToken
		return nil
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2015 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"bytes"
	"fmt"
	"io"
	"unicode/utf8"
)

// MessageHeaderSize is the number of bytes in a bitcoin message header.
// Bitcoin network (magic) 4 bytes + command 12 bytes + payload length 4 bytes +
// checksum 4 bytes.
const MessageHeaderSize = 24

// CommandSize is the fixed size of all commands in the common bitcoin message
// header.  Shorter commands must be zero padded.
const CommandSize = 12

// MaxMessagePayload is the maximum bytes a message can be regardless of other
// individual limits imposed by messages themselves.
const MaxMessagePayload = (1024 * 1024 * 32) // 32MB

// Commands used in bitcoin message headers which describe the type of message.
const (
	CmdVersion     = "version"
	CmdVerAck      = "verack"
	CmdGetAddr     = "getaddr"
	CmdAddr        = "addr"
	CmdGetBlocks   = "getblocks"
	CmdInv         = "inv"
	CmdGetData     = "getdata"
	CmdNotFound    = "notfound"
	CmdBlock       = "block"     // Factoid-0 Block
	CmdTx          = "tx"        // Factoid-0 Tx
	CmdFBlock      = "FBlock"    // Factoid-1 Block
	CmdFactoidTX   = "factoidtx" // Factoid-1 Tx
	CmdGetHeaders  = "getheaders"
	CmdHeaders     = "headers"
	CmdPing        = "ping"
	CmdPong        = "pong"
	CmdAlert       = "alert"
	CmdMemPool     = "mempool"
	CmdFilterAdd   = "filteradd"
	CmdFilterClear = "filterclear"
	CmdFilterLoad  = "filterload"
	CmdReject      = "reject"
	// TODO remove CmdTestCred before production
	CmdTestCredit = "testcredit"

	// Factom additions:

	// Factom downloads: block & entry
	CmdGetDirBlocks  = "getdirblocks"
	CmdDirInv        = "dirInv"
	CmdGetDirData    = "getdirdata"
	CmdDirBlock      = "dirblock"
	CmdGetNonDirData = "getnodirdata"
	CmdABlock        = "adminblock" // Admin Block
	CmdECBlock       = "ecblock"    // Entry Credit Block
	CmdEBlock        = "entryblock" // Entry Block
	CmdGetEntryData  = "getentrydata"
	CmdEntry         = "entry"

	CmdCommitChain = "commitchain"
	CmdRevealChain = "revealchain"
	CmdCommitEntry = "commitentry"
	CmdRevealEntry = "revealentry"

	CmdAcknowledgement = "confirmation"
	CmdMHashReveal     = "mhashreveal"
)

// MaxAppMsgPayload is the maximum bytes a factom app message can be in bytes.
const MaxAppMsgPayload = (10 * 1024) // 10Kib

// Message is an interface that describes a bitcoin message.  A type that
// implements Message has complete control over the representation of its data
// and may therefore contain additional or fewer fields than those which
// are used directly in the protocol encoded message.
type Message interface {
	BtcDecode(io.Reader, uint32) error
	BtcEncode(io.Writer, uint32) error
	Command() string
	MaxPayloadLength(uint32) uint32
}

// makeEmptyMessage creates a message of the appropriate concrete type based
// on the command.
func makeEmptyMessage(command string) (Message, error) {
	var msg Message
	switch command {
	case CmdVersion:
		msg = &MsgVersion{}

	case CmdVerAck:
		msg = &MsgVerAck{}

	case CmdFBlock:
		msg = &MsgFBlock{}

	case CmdGetAddr:
		msg = &MsgGetAddr{}

	case CmdAddr:
		msg = &MsgAddr{}

	case CmdGetBlocks:
		msg = &MsgGetBlocks{}

	case CmdBlock:
		msg = &MsgBlock{}

	case CmdInv:
		msg = &MsgInv{}

	case CmdGetData:
		msg = &MsgGetData{}

	case CmdNotFound:
		msg = &MsgNotFound{}

	case CmdTx:
		msg = &MsgTx{}

	case CmdPing:
		msg = &MsgPing{}

	case CmdPong:
		msg = &MsgPong{}

	case CmdGetHeaders:
		msg = &MsgGetHeaders{}

	case CmdHeaders:
		msg = &MsgHeaders{}

	case CmdAlert:
		msg = &MsgAlert{}

	case CmdMemPool:
		msg = &MsgMemPool{}

	case CmdFilterAdd:
		msg = &MsgFilterAdd{}

	case CmdFilterClear:
		msg = &MsgFilterClear{}

	case CmdFilterLoad:
		msg = &MsgFilterLoad{}

	case CmdReject:
		msg = &MsgReject{}

	case CmdRevealEntry:
		msg = &MsgRevealEntry{}

	case CmdCommitEntry:
		msg = &MsgCommitEntry{}

	case CmdRevealChain:
		msg = &MsgRevealChain{}

	case CmdCommitChain:
		msg = &MsgCommitChain{}

	case CmdAcknowledgement:
		msg = &MsgAcknowledgement{}

	case CmdGetDirBlocks:
		msg = &MsgGetDirBlocks{}

	case CmdDirInv:
		msg = &MsgDirInv{}

	case CmdGetDirData:
		msg = &MsgGetDirData{}

	case CmdDirBlock:
		msg = &MsgDirBlock{}

	case CmdGetNonDirData:
		msg = &MsgGetNonDirData{}

	case CmdABlock:
		msg = &MsgABlock{}

	case CmdECBlock:
		msg = &MsgECBlock{}

	case CmdEBlock:
		msg = &MsgEBlock{}

	case CmdGetEntryData:
		msg = &MsgGetEntryData{}

	case CmdEntry:
		msg = &MsgEntry{}

	default:
		return nil, fmt.Errorf("unhandled command [%s]", command)
	}
	return msg, nil
}

// messageHeader defines the header structure for all bitcoin protocol messages.
type messageHeader struct {
	magic    BitcoinNet // 4 bytes
	command  string     // 12 bytes
	length   uint32     // 4 bytes
	checksum [4]byte    // 4 bytes
}

// readMessageHeader reads a bitcoin message header from r.
func readMessageHeader(r io.Reader) (int, *messageHeader, error) {
	// Since readElements doesn't return the amount of bytes read, attempt
	// to read the entire header into a buffer first in case there is a
	// short read so the proper amount of read bytes are known.  This works
	// since the header is a fixed size.
	var headerBytes [MessageHeaderSize]byte
	n, err := io.ReadFull(r, headerBytes[:])
	if err != nil {
		return n, nil, err
	}
	hr := bytes.NewReader(headerBytes[:])

	// Create and populate a messageHeader struct from the raw header bytes.
	hdr := messageHeader{}
	var command [CommandSize]byte
	readElements(hr, &hdr.magic, &command, &hdr.length, &hdr.checksum)

	// Strip trailing zeros from command string.
	hdr.command = string(bytes.TrimRight(command[:], string(0)))

	return n, &hdr, nil
}

// discardInput reads n bytes from reader r in chunks and discards the read
// bytes.  This is used to skip payloads when various errors occur and helps
// prevent rogue nodes from causing massive memory allocation through forging
// header length.
func discardInput(r io.Reader, n uint32) {
	maxSize := uint32(10 * 1024) // 10k at a time
	numReads := n / maxSize
	bytesRemaining := n % maxSize
	if n > 0 {
		buf := make([]byte, maxSize)
		for i := uint32(0); i < numReads; i++ {
			io.ReadFull(r, buf)
		}
	}
	if bytesRemaining > 0 {
		buf := make([]byte, bytesRemaining)
		io.ReadFull(r, buf)
	}
}

// WriteMessageN writes a bitcoin Message to w including the necessary header
// information and returns the number of bytes written.    This function is the
// same as WriteMessage except it also returns the number of bytes written.
func WriteMessageN(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) (int, error) {
	totalBytes := 0

	// Enforce max command size.
	var command [CommandSize]byte
	cmd := msg.Command()
	if len(cmd) > CommandSize {
		str := fmt.Sprintf("command [%s] is too long [max %v]",
			cmd, CommandSize)
		return totalBytes, messageError("WriteMessage", str)
	}
	copy(command[:], []byte(cmd))

	// Encode the message payload.
	var bw bytes.Buffer
	err := msg.BtcEncode(&bw, pver)
	if err != nil {
		return totalBytes, err
	}
	payload := bw.Bytes()
	lenp := len(payload)

	// Enforce maximum overall message payload.
	if lenp > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload is %d bytes",
			lenp, MaxMessagePayload)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Enforce maximum message payload based on the message type.
	mpl := msg.MaxPayloadLength(pver)
	if uint32(lenp) > mpl {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload size for "+
			"messages of type [%s] is %d.", lenp, cmd, mpl)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Create header for the message.
	hdr := messageHeader{}
	hdr.magic = btcnet
	hdr.command = cmd
	hdr.length = uint32(lenp)
	copy(hdr.checksum[:], DoubleSha256(payload)[0:4])

	// Encode the header for the message.  This is done to a buffer
	// rather than directly to the writer since writeElements doesn't
	// return the number of bytes written.
	hw := bytes.NewBuffer(make([]byte, 0, MessageHeaderSize))
	writeElements(hw, hdr.magic, command, hdr.length, hdr.checksum)

	// Write header.
	n, err := w.Write(hw.Bytes())
	if err != nil {
		totalBytes += n
		return totalBytes, err
	}
	totalBytes += n

	// Write payload.
	n, err = w.Write(payload)
	if err != nil {
		totalBytes += n
		return totalBytes, err
	}
	totalBytes += n

	return totalBytes, nil
}

// WriteMessage writes a bitcoin Message to w including the necessary header
// information.  This function is the same as WriteMessageN except it doesn't
// doesn't return the number of bytes written.  This function is mainly provided
// for backwards compatibility with the original API, but it's also useful for
// callers that don't care about byte counts.
func WriteMessage(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) error {
	_, err := WriteMessageN(w, msg, pver, btcnet)
	return err
}

// ReadMessageN reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the number of
// bytes read in addition to the parsed Message and raw bytes which comprise the
// message.  This function is the same as ReadMessage except it also returns the
// number of bytes read.
func ReadMessageN(r io.Reader, pver uint32, btcnet BitcoinNet) (int, Message, []byte, error) {
	totalBytes := 0
	n, hdr, err := readMessageHeader(r)
	if err != nil {
		totalBytes += n
		return totalBytes, nil, nil, err
	}
	totalBytes += n

	// Enforce maximum message payload.
	if hdr.length > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - header "+
			"indicates %d bytes, but max message payload is %d "+
			"bytes.", hdr.length, MaxMessagePayload)
		return totalBytes, nil, nil, messageError("ReadMessage", str)

	}

	// Check for messages from the wrong bitcoin network.
	if hdr.magic != btcnet {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("message from other network [%v]", hdr.magic)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Check for malformed commands.
	command := hdr.command
	if !utf8.ValidString(command) {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("invalid command %v", []byte(command))
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Create struct of appropriate message type based on the command.
	msg, err := makeEmptyMessage(command)
	if err != nil {
		discardInput(r, hdr.length)
		return totalBytes, nil, nil, messageError("ReadMessage",
			err.Error())
	}

	// Check for maximum length based on the message type as a malicious client
	// could otherwise create a well-formed header and set the length to max
	// numbers in order to exhaust the machine's memory.
	mpl := msg.MaxPayloadLength(pver)
	if hdr.length > mpl {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("payload exceeds max length - header "+
			"indicates %v bytes, but max payload size for "+
			"messages of type [%v] is %v.", hdr.length, command, mpl)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Read payload.
	payload := make([]byte, hdr.length)
	n, err = io.ReadFull(r, payload)
	if err != nil {
		totalBytes += n
		return totalBytes, nil, nil, err
	}
	totalBytes += n

	// Test checksum.
	checksum := DoubleSha256(payload)[0:4]
	if !bytes.Equal(checksum[:], hdr.checksum[:]) {
		str := fmt.Sprintf("payload checksum failed - header "+
			"indicates %v, but actual checksum is %v.",
			hdr.checksum, checksum)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Unmarshal message.  NOTE: This must be a *bytes.Buffer since the
	// MsgVersion BtcDecode function requires it.
	pr := bytes.NewBuffer(payload)
	err = msg.BtcDecode(pr, pver)
	if err != nil {
		return totalBytes, nil, nil, err
	}

	return totalBytes, msg, payload, nil
}

// ReadMessage reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the parsed
// Message and raw bytes which comprise the message.  This function only differs
// from ReadMessageN in that it doesn't return the number of bytes read.  This
// function is mainly provided for backwards compatibility with the original
// API, but it's also useful for callers that don't care about byte counts.
func ReadMessage(r io.Reader, pver uint32, btcnet BitcoinNet) (Message, []byte, error) {
	_, msg, buf, err := ReadMessageN(r, pver, btcnet)
	return msg, buf, err
}
<commit_msg>factoidtx reception fix<commit_after>// Copyright (c) 2013-2015 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"bytes"
	"fmt"
	"io"
	"unicode/utf8"
)

// MessageHeaderSize is the number of bytes in a bitcoin message header.
// Bitcoin network (magic) 4 bytes + command 12 bytes + payload length 4 bytes +
// checksum 4 bytes.
const MessageHeaderSize = 24

// CommandSize is the fixed size of all commands in the common bitcoin message
// header.  Shorter commands must be zero padded.
const CommandSize = 12

// MaxMessagePayload is the maximum bytes a message can be regardless of other
// individual limits imposed by messages themselves.
const MaxMessagePayload = (1024 * 1024 * 32) // 32MB

// Commands used in bitcoin message headers which describe the type of message.
const (
	CmdVersion     = "version"
	CmdVerAck      = "verack"
	CmdGetAddr     = "getaddr"
	CmdAddr        = "addr"
	CmdGetBlocks   = "getblocks"
	CmdInv         = "inv"
	CmdGetData     = "getdata"
	CmdNotFound    = "notfound"
	CmdBlock       = "block"     // Factoid-0 Block
	CmdTx          = "tx"        // Factoid-0 Tx
	CmdFBlock      = "FBlock"    // Factoid-1 Block
	CmdFactoidTX   = "factoidtx" // Factoid-1 Tx
	CmdGetHeaders  = "getheaders"
	CmdHeaders     = "headers"
	CmdPing        = "ping"
	CmdPong        = "pong"
	CmdAlert       = "alert"
	CmdMemPool     = "mempool"
	CmdFilterAdd   = "filteradd"
	CmdFilterClear = "filterclear"
	CmdFilterLoad  = "filterload"
	CmdReject      = "reject"
	// TODO remove CmdTestCred before production
	CmdTestCredit = "testcredit"

	// Factom additions:

	// Factom downloads: block & entry
	CmdGetDirBlocks  = "getdirblocks"
	CmdDirInv        = "dirInv"
	CmdGetDirData    = "getdirdata"
	CmdDirBlock      = "dirblock"
	CmdGetNonDirData = "getnodirdata"
	CmdABlock        = "adminblock" // Admin Block
	CmdECBlock       = "ecblock"    // Entry Credit Block
	CmdEBlock        = "entryblock" // Entry Block
	CmdGetEntryData  = "getentrydata"
	CmdEntry         = "entry"

	CmdCommitChain = "commitchain"
	CmdRevealChain = "revealchain"
	CmdCommitEntry = "commitentry"
	CmdRevealEntry = "revealentry"

	CmdAcknowledgement = "confirmation"
	CmdMHashReveal     = "mhashreveal"
)

// MaxAppMsgPayload is the maximum bytes a factom app message can be in bytes.
const MaxAppMsgPayload = (10 * 1024) // 10Kib

// Message is an interface that describes a bitcoin message.  A type that
// implements Message has complete control over the representation of its data
// and may therefore contain additional or fewer fields than those which
// are used directly in the protocol encoded message.
type Message interface {
	BtcDecode(io.Reader, uint32) error
	BtcEncode(io.Writer, uint32) error
	Command() string
	MaxPayloadLength(uint32) uint32
}

// makeEmptyMessage creates a message of the appropriate concrete type based
// on the command.
func makeEmptyMessage(command string) (Message, error) {
	var msg Message
	switch command {
	case CmdVersion:
		msg = &MsgVersion{}

	case CmdVerAck:
		msg = &MsgVerAck{}

	case CmdFBlock:
		msg = &MsgFBlock{}

	case CmdGetAddr:
		msg = &MsgGetAddr{}

	case CmdAddr:
		msg = &MsgAddr{}

	case CmdGetBlocks:
		msg = &MsgGetBlocks{}

	case CmdBlock:
		msg = &MsgBlock{}

	case CmdInv:
		msg = &MsgInv{}

	case CmdGetData:
		msg = &MsgGetData{}

	case CmdNotFound:
		msg = &MsgNotFound{}

		/*
			case CmdTx:
				msg = &MsgTx{}
		*/

	case CmdPing:
		msg = &MsgPing{}

	case CmdPong:
		msg = &MsgPong{}

	case CmdGetHeaders:
		msg = &MsgGetHeaders{}

	case CmdHeaders:
		msg = &MsgHeaders{}

	case CmdAlert:
		msg = &MsgAlert{}

	case CmdMemPool:
		msg = &MsgMemPool{}

	case CmdFilterAdd:
		msg = &MsgFilterAdd{}

	case CmdFilterClear:
		msg = &MsgFilterClear{}

	case CmdFilterLoad:
		msg = &MsgFilterLoad{}

	case CmdReject:
		msg = &MsgReject{}

	case CmdRevealEntry:
		msg = &MsgRevealEntry{}

	case CmdCommitEntry:
		msg = &MsgCommitEntry{}

	case CmdRevealChain:
		msg = &MsgRevealChain{}

	case CmdCommitChain:
		msg = &MsgCommitChain{}

	case CmdAcknowledgement:
		msg = &MsgAcknowledgement{}

	case CmdGetDirBlocks:
		msg = &MsgGetDirBlocks{}

	case CmdDirInv:
		msg = &MsgDirInv{}

	case CmdGetDirData:
		msg = &MsgGetDirData{}

	case CmdDirBlock:
		msg = &MsgDirBlock{}

	case CmdGetNonDirData:
		msg = &MsgGetNonDirData{}

	case CmdABlock:
		msg = &MsgABlock{}

	case CmdECBlock:
		msg = &MsgECBlock{}

	case CmdEBlock:
		msg = &MsgEBlock{}

	case CmdGetEntryData:
		msg = &MsgGetEntryData{}

	case CmdEntry:
		msg = &MsgEntry{}

	case CmdFactoidTX:
		msg = &MsgFactoidTX{}

	default:
		return nil, fmt.Errorf("unhandled command [%s]", command)
	}
	return msg, nil
}

// messageHeader defines the header structure for all bitcoin protocol messages.
type messageHeader struct {
	magic    BitcoinNet // 4 bytes
	command  string     // 12 bytes
	length   uint32     // 4 bytes
	checksum [4]byte    // 4 bytes
}

// readMessageHeader reads a bitcoin message header from r.
func readMessageHeader(r io.Reader) (int, *messageHeader, error) {
	// Since readElements doesn't return the amount of bytes read, attempt
	// to read the entire header into a buffer first in case there is a
	// short read so the proper amount of read bytes are known.  This works
	// since the header is a fixed size.
	var headerBytes [MessageHeaderSize]byte
	n, err := io.ReadFull(r, headerBytes[:])
	if err != nil {
		return n, nil, err
	}
	hr := bytes.NewReader(headerBytes[:])

	// Create and populate a messageHeader struct from the raw header bytes.
	hdr := messageHeader{}
	var command [CommandSize]byte
	readElements(hr, &hdr.magic, &command, &hdr.length, &hdr.checksum)

	// Strip trailing zeros from command string.
	hdr.command = string(bytes.TrimRight(command[:], string(0)))

	return n, &hdr, nil
}

// discardInput reads n bytes from reader r in chunks and discards the read
// bytes.  This is used to skip payloads when various errors occur and helps
// prevent rogue nodes from causing massive memory allocation through forging
// header length.
func discardInput(r io.Reader, n uint32) {
	maxSize := uint32(10 * 1024) // 10k at a time
	numReads := n / maxSize
	bytesRemaining := n % maxSize
	if n > 0 {
		buf := make([]byte, maxSize)
		for i := uint32(0); i < numReads; i++ {
			io.ReadFull(r, buf)
		}
	}
	if bytesRemaining > 0 {
		buf := make([]byte, bytesRemaining)
		io.ReadFull(r, buf)
	}
}

// WriteMessageN writes a bitcoin Message to w including the necessary header
// information and returns the number of bytes written.    This function is the
// same as WriteMessage except it also returns the number of bytes written.
func WriteMessageN(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) (int, error) {
	totalBytes := 0

	// Enforce max command size.
	var command [CommandSize]byte
	cmd := msg.Command()
	if len(cmd) > CommandSize {
		str := fmt.Sprintf("command [%s] is too long [max %v]",
			cmd, CommandSize)
		return totalBytes, messageError("WriteMessage", str)
	}
	copy(command[:], []byte(cmd))

	// Encode the message payload.
	var bw bytes.Buffer
	err := msg.BtcEncode(&bw, pver)
	if err != nil {
		return totalBytes, err
	}
	payload := bw.Bytes()
	lenp := len(payload)

	// Enforce maximum overall message payload.
	if lenp > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload is %d bytes",
			lenp, MaxMessagePayload)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Enforce maximum message payload based on the message type.
	mpl := msg.MaxPayloadLength(pver)
	if uint32(lenp) > mpl {
		str := fmt.Sprintf("message payload is too large - encoded "+
			"%d bytes, but maximum message payload size for "+
			"messages of type [%s] is %d.", lenp, cmd, mpl)
		return totalBytes, messageError("WriteMessage", str)
	}

	// Create header for the message.
	hdr := messageHeader{}
	hdr.magic = btcnet
	hdr.command = cmd
	hdr.length = uint32(lenp)
	copy(hdr.checksum[:], DoubleSha256(payload)[0:4])

	// Encode the header for the message.  This is done to a buffer
	// rather than directly to the writer since writeElements doesn't
	// return the number of bytes written.
	hw := bytes.NewBuffer(make([]byte, 0, MessageHeaderSize))
	writeElements(hw, hdr.magic, command, hdr.length, hdr.checksum)

	// Write header.
	n, err := w.Write(hw.Bytes())
	if err != nil {
		totalBytes += n
		return totalBytes, err
	}
	totalBytes += n

	// Write payload.
	n, err = w.Write(payload)
	if err != nil {
		totalBytes += n
		return totalBytes, err
	}
	totalBytes += n

	return totalBytes, nil
}

// WriteMessage writes a bitcoin Message to w including the necessary header
// information.  This function is the same as WriteMessageN except it doesn't
// doesn't return the number of bytes written.  This function is mainly provided
// for backwards compatibility with the original API, but it's also useful for
// callers that don't care about byte counts.
func WriteMessage(w io.Writer, msg Message, pver uint32, btcnet BitcoinNet) error {
	_, err := WriteMessageN(w, msg, pver, btcnet)
	return err
}

// ReadMessageN reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the number of
// bytes read in addition to the parsed Message and raw bytes which comprise the
// message.  This function is the same as ReadMessage except it also returns the
// number of bytes read.
func ReadMessageN(r io.Reader, pver uint32, btcnet BitcoinNet) (int, Message, []byte, error) {
	totalBytes := 0
	n, hdr, err := readMessageHeader(r)
	if err != nil {
		totalBytes += n
		return totalBytes, nil, nil, err
	}
	totalBytes += n

	// Enforce maximum message payload.
	if hdr.length > MaxMessagePayload {
		str := fmt.Sprintf("message payload is too large - header "+
			"indicates %d bytes, but max message payload is %d "+
			"bytes.", hdr.length, MaxMessagePayload)
		return totalBytes, nil, nil, messageError("ReadMessage", str)

	}

	// Check for messages from the wrong bitcoin network.
	if hdr.magic != btcnet {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("message from other network [%v]", hdr.magic)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Check for malformed commands.
	command := hdr.command
	if !utf8.ValidString(command) {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("invalid command %v", []byte(command))
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Create struct of appropriate message type based on the command.
	msg, err := makeEmptyMessage(command)
	if err != nil {
		discardInput(r, hdr.length)
		return totalBytes, nil, nil, messageError("ReadMessage",
			err.Error())
	}

	// Check for maximum length based on the message type as a malicious client
	// could otherwise create a well-formed header and set the length to max
	// numbers in order to exhaust the machine's memory.
	mpl := msg.MaxPayloadLength(pver)
	if hdr.length > mpl {
		discardInput(r, hdr.length)
		str := fmt.Sprintf("payload exceeds max length - header "+
			"indicates %v bytes, but max payload size for "+
			"messages of type [%v] is %v.", hdr.length, command, mpl)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Read payload.
	payload := make([]byte, hdr.length)
	n, err = io.ReadFull(r, payload)
	if err != nil {
		totalBytes += n
		return totalBytes, nil, nil, err
	}
	totalBytes += n

	// Test checksum.
	checksum := DoubleSha256(payload)[0:4]
	if !bytes.Equal(checksum[:], hdr.checksum[:]) {
		str := fmt.Sprintf("payload checksum failed - header "+
			"indicates %v, but actual checksum is %v.",
			hdr.checksum, checksum)
		return totalBytes, nil, nil, messageError("ReadMessage", str)
	}

	// Unmarshal message.  NOTE: This must be a *bytes.Buffer since the
	// MsgVersion BtcDecode function requires it.
	pr := bytes.NewBuffer(payload)
	err = msg.BtcDecode(pr, pver)
	if err != nil {
		return totalBytes, nil, nil, err
	}

	return totalBytes, msg, payload, nil
}

// ReadMessage reads, validates, and parses the next bitcoin Message from r for
// the provided protocol version and bitcoin network.  It returns the parsed
// Message and raw bytes which comprise the message.  This function only differs
// from ReadMessageN in that it doesn't return the number of bytes read.  This
// function is mainly provided for backwards compatibility with the original
// API, but it's also useful for callers that don't care about byte counts.
func ReadMessage(r io.Reader, pver uint32, btcnet BitcoinNet) (Message, []byte, error) {
	_, msg, buf, err := ReadMessageN(r, pver, btcnet)
	return msg, buf, err
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"io"
	"os"
	"strconv"
	"strings"
	"testing"
)

var testWriter io.Writer = os.Stdout

// RandomPasswordApp([]string{}) should print an error about not enough args
// codegangsta's cli currently throws a runtime error if run is called with an empty slice
// Not sure if this scenario is even possible in practice
/*
func TestRandomPasswordNoArgs(t *testing.T) {
	buf := &bytes.Buffer{}
	testWriter = buf

	app := randomPasswordApp(testWriter)
	app.Run([]string{})

	message := "Called with no executable name or other arguments. This function needs an executable name."
	messageln := message + "\n"

	if buf.String() != messageln {
		t.Error("randomPassword([]string{}) printed '", buf.String(), "' expected '", messageln, "'")
	}
}
*/

// RandomPasswordApp([]string{passgen}) should exhibit default app behavior
// That is, chars=a1! and length=20
func TestRandomPasswordOneArg(t *testing.T) {
	premessage := "Random password: "
	postmessage1 := " - "
	postmessage2 := " characters long\n"

	buf := &bytes.Buffer{}
	testWriter = buf
	app := randomPasswordApp(testWriter)

	input := []string{"passgen"}
	app.Run(input)

	n := 20
	pre := buf.String()[:len(premessage)]
	pass := buf.String()[len(premessage) : len(premessage)+n]
	post1 := buf.String()[len(premessage)+n : len(premessage)+n+len(postmessage1)]
	_ = buf.String()[len(premessage)+n+len(postmessage1) : len(premessage)+n+len(postmessage1)+len(strconv.Itoa(n))]
	post2 := buf.String()[len(premessage)+n+len(postmessage1)+len(strconv.Itoa(n)):]

	if pre != premessage {
		t.Error("Expected premessage of: '" + premessage + "' Got: '" + pre + "'")
	}
	if post1 != postmessage1 {
		t.Error("Expected postmessage of: '" + postmessage1 + "' Got: '" + post1 + "'")
	}
	if post2 != postmessage2 {
		t.Error("Expected postmessage of: '" + postmessage2 + "' Got: '" + post2 + "'")
	}

	// Check char set membership
	all := union(alphabet, digit, special)
	for _, x := range []rune(pass) {
		if !all[string(x)] {
			t.Error("pass =", pass, "contains", string(x), "but this is not in", all)
		}
	}

}

// RandomPasswordApp([]string{passgen, n}) n is not a number should print error about atoi or invalid length arg
func TestRandomPasswordNAN(t *testing.T) {
	buf := &bytes.Buffer{}
	testWriter = buf
	n := "stuff"

	app := randomPasswordApp(testWriter)
	app.Run([]string{"passgen", "-l", n})

	expected := `Incorrect Usage.

NAME:
   passgen - generate a random password

USAGE:
   /tmp/go-build496363236/github.com/voutasaurus/passgen/_test/passgen.test [global options] command [command options] [arguments...]

VERSION:
   0.0.0

COMMANDS:
   help, h	Shows a list of commands or help for one command

GLOBAL OPTIONS:
   --chars, -c "a1*"	specify which character sets to include in the generated password
   --length, -l "20"	specify the length of the generated password
   --help, -h		show help
   --version, -v	print the version

`
	expectWords := strings.Fields(expected)
	actualWords := strings.Fields(buf.String())

	if len(actualWords) != len(expectWords) {
		t.Error("expected", len(expectWords), "words. got", len(actualWords), "words.")
		t.Error("passgen -l", n, "printed '", buf.String(), "' expected", expected)
		t.FailNow()
	}

	for i := range actualWords {
		if i != 10 && actualWords[i] != expectWords[i] { // i == 10 contains the build number for the tests
			t.Error("Output at word", i, "expected:", expectWords[i], "got:", actualWords[i])
			t.Error("passgen -l", n, "printed '", buf.String(), "' expected", expected)
		}
	}

}

// RandomPasswordApp([]string{passgen, -1})
func TestRandomPasswordNegLength(t *testing.T) {
	buf := &bytes.Buffer{}
	testWriter = buf
	n := "-1"

	app := randomPasswordApp(testWriter)
	app.Run([]string{"passgen", "-l", n})

	if !strings.Contains(buf.String(), "is not a valid password length. Must be greater than zero.") {
		t.Error("passgen -l", n, "printed '", buf.String(), "' expected invalid password length message")
	}
}

// RandomPasswordApp([]string{passgen, n}) where n>0 should print n random characters with a message
func TestRandomPasswordJustN(t *testing.T) {

	premessage := "Random password: "
	postmessage1 := " - "
	postmessage2 := " characters long\n"

	buf := &bytes.Buffer{}
	testWriter = buf
	app := randomPasswordApp(testWriter)

	for n := 0; n < 100; n++ {
		input := []string{"passgen", "-l", strconv.Itoa(n)}
		app.Run(input)

		pre := buf.String()[:len(premessage)]
		pass := buf.String()[len(premessage) : len(premessage)+n]
		post1 := buf.String()[len(premessage)+n : len(premessage)+n+len(postmessage1)]
		_ = buf.String()[len(premessage)+n+len(postmessage1) : len(premessage)+n+len(postmessage1)+len(strconv.Itoa(n))]
		post2 := buf.String()[len(premessage)+n+len(postmessage1)+len(strconv.Itoa(n)):]

		if pre != premessage {
			t.Error("Expected premessage of: '" + premessage + "' Got: '" + pre + "'")
		}
		if pass == "" && n != 0 {
			t.Error("passgen 0 should return empty password.")
		}
		if post1 != postmessage1 {
			t.Error("Expected postmessage of: '" + postmessage1 + "' Got: '" + post1 + "'")
		}
		if post2 != postmessage2 {
			t.Error("Expected postmessage of: '" + postmessage2 + "' Got: '" + post2 + "'")
		}

		// reset buffer for next test case
		buf.Reset()
	}
}

// RandomPasswordApp([]string{"passgen", n, set}) should print characters from set
func TestRandomPasswordSetMembership(t *testing.T) {

	premessage := "Random password: "

	buf := &bytes.Buffer{}
	testWriter = buf

	app := randomPasswordApp(testWriter)

	all := union(alphabet, digit, special)

	input := []string{"passgen", "-l", strconv.Itoa(0)}
	for n := 0; n < 100; n++ {
		input[2] = strconv.Itoa(n)
		app.Run(input)

		pass := buf.String()[len(premessage) : len(premessage)+n]

		for _, x := range []rune(pass) {
			if !all[string(x)] {
				t.Error("pass =", pass, "contains", string(x), "but this is not in", all)
			}
		}

		// reset buffer for next test case
		buf.Reset()
	}
}
<commit_msg>Updates help output in test case<commit_after>package main

import (
	"bytes"
	"io"
	"os"
	"strconv"
	"strings"
	"testing"
)

var testWriter io.Writer = os.Stdout

// RandomPasswordApp([]string{}) should print an error about not enough args
// codegangsta's cli currently throws a runtime error if run is called with an empty slice
// Not sure if this scenario is even possible in practice
/*
func TestRandomPasswordNoArgs(t *testing.T) {
	buf := &bytes.Buffer{}
	testWriter = buf

	app := randomPasswordApp(testWriter)
	app.Run([]string{})

	message := "Called with no executable name or other arguments. This function needs an executable name."
	messageln := message + "\n"

	if buf.String() != messageln {
		t.Error("randomPassword([]string{}) printed '", buf.String(), "' expected '", messageln, "'")
	}
}
*/

// RandomPasswordApp([]string{passgen}) should exhibit default app behavior
// That is, chars=a1! and length=20
func TestRandomPasswordOneArg(t *testing.T) {
	premessage := "Random password: "
	postmessage1 := " - "
	postmessage2 := " characters long\n"

	buf := &bytes.Buffer{}
	testWriter = buf
	app := randomPasswordApp(testWriter)

	input := []string{"passgen"}
	app.Run(input)

	n := 20
	pre := buf.String()[:len(premessage)]
	pass := buf.String()[len(premessage) : len(premessage)+n]
	post1 := buf.String()[len(premessage)+n : len(premessage)+n+len(postmessage1)]
	_ = buf.String()[len(premessage)+n+len(postmessage1) : len(premessage)+n+len(postmessage1)+len(strconv.Itoa(n))]
	post2 := buf.String()[len(premessage)+n+len(postmessage1)+len(strconv.Itoa(n)):]

	if pre != premessage {
		t.Error("Expected premessage of: '" + premessage + "' Got: '" + pre + "'")
	}
	if post1 != postmessage1 {
		t.Error("Expected postmessage of: '" + postmessage1 + "' Got: '" + post1 + "'")
	}
	if post2 != postmessage2 {
		t.Error("Expected postmessage of: '" + postmessage2 + "' Got: '" + post2 + "'")
	}

	// Check char set membership
	all := union(alphabet, digit, special)
	for _, x := range []rune(pass) {
		if !all[string(x)] {
			t.Error("pass =", pass, "contains", string(x), "but this is not in", all)
		}
	}

}

// RandomPasswordApp([]string{passgen, n}) n is not a number should print error about atoi or invalid length arg
func TestRandomPasswordNAN(t *testing.T) {
	buf := &bytes.Buffer{}
	testWriter = buf
	n := "stuff"

	app := randomPasswordApp(testWriter)
	app.Run([]string{"passgen", "-l", n})

	expected := `Incorrect Usage.

NAME:
   passgen - generate a random password

USAGE:
   /tmp/go-build496363236/github.com/voutasaurus/passgen/_test/passgen.test [global options] command [command options] [arguments...]

VERSION:
   0.0.0

COMMANDS:
   help, h	Shows a list of commands or help for one command

GLOBAL OPTIONS:
   --chars, -c "a1*"	specify which character sets to include in the generated password
   --length, -l "20"	specify the length of the generated password
   --words, -w          use words instead
   --help, -h		show help
   --version, -v	print the version

`
	expectWords := strings.Fields(expected)
	actualWords := strings.Fields(buf.String())

	if len(actualWords) != len(expectWords) {
		t.Error("expected", len(expectWords), "words. got", len(actualWords), "words.")
		t.Error("passgen -l", n, "printed '", buf.String(), "' expected", expected)
		t.FailNow()
	}

	for i := range actualWords {
		if i != 10 && actualWords[i] != expectWords[i] { // i == 10 contains the build number for the tests
			t.Error("Output at word", i, "expected:", expectWords[i], "got:", actualWords[i])
			t.Error("passgen -l", n, "printed '", buf.String(), "' expected", expected)
		}
	}

}

// RandomPasswordApp([]string{passgen, -1})
func TestRandomPasswordNegLength(t *testing.T) {
	buf := &bytes.Buffer{}
	testWriter = buf
	n := "-1"

	app := randomPasswordApp(testWriter)
	app.Run([]string{"passgen", "-l", n})

	if !strings.Contains(buf.String(), "is not a valid password length. Must be greater than zero.") {
		t.Error("passgen -l", n, "printed '", buf.String(), "' expected invalid password length message")
	}
}

// RandomPasswordApp([]string{passgen, n}) where n>0 should print n random characters with a message
func TestRandomPasswordJustN(t *testing.T) {

	premessage := "Random password: "
	postmessage1 := " - "
	postmessage2 := " characters long\n"

	buf := &bytes.Buffer{}
	testWriter = buf
	app := randomPasswordApp(testWriter)

	for n := 0; n < 100; n++ {
		input := []string{"passgen", "-l", strconv.Itoa(n)}
		app.Run(input)

		pre := buf.String()[:len(premessage)]
		pass := buf.String()[len(premessage) : len(premessage)+n]
		post1 := buf.String()[len(premessage)+n : len(premessage)+n+len(postmessage1)]
		_ = buf.String()[len(premessage)+n+len(postmessage1) : len(premessage)+n+len(postmessage1)+len(strconv.Itoa(n))]
		post2 := buf.String()[len(premessage)+n+len(postmessage1)+len(strconv.Itoa(n)):]

		if pre != premessage {
			t.Error("Expected premessage of: '" + premessage + "' Got: '" + pre + "'")
		}
		if pass == "" && n != 0 {
			t.Error("passgen 0 should return empty password.")
		}
		if post1 != postmessage1 {
			t.Error("Expected postmessage of: '" + postmessage1 + "' Got: '" + post1 + "'")
		}
		if post2 != postmessage2 {
			t.Error("Expected postmessage of: '" + postmessage2 + "' Got: '" + post2 + "'")
		}

		// reset buffer for next test case
		buf.Reset()
	}
}

// RandomPasswordApp([]string{"passgen", n, set}) should print characters from set
func TestRandomPasswordSetMembership(t *testing.T) {

	premessage := "Random password: "

	buf := &bytes.Buffer{}
	testWriter = buf

	app := randomPasswordApp(testWriter)

	all := union(alphabet, digit, special)

	input := []string{"passgen", "-l", strconv.Itoa(0)}
	for n := 0; n < 100; n++ {
		input[2] = strconv.Itoa(n)
		app.Run(input)

		pass := buf.String()[len(premessage) : len(premessage)+n]

		for _, x := range []rune(pass) {
			if !all[string(x)] {
				t.Error("pass =", pass, "contains", string(x), "but this is not in", all)
			}
		}

		// reset buffer for next test case
		buf.Reset()
	}
}
<|endoftext|>