<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/gophergala2016/goad/queue"
)

func main() {
	address := os.Args[1]
	concurrencycount, err := strconv.Atoi(os.Args[2])
	if err != nil {
		fmt.Printf("ERROR %s\n", err)
		return
	}
	maxRequestCount, err := strconv.Atoi(os.Args[3])
	sqsurl := os.Args[4]
	awsregion := os.Args[5]
	clientTimeout, _ := time.ParseDuration("1s")
	if len(os.Args) > 6 {
		newClientTimeout, err := time.ParseDuration(os.Args[6])
		if err == nil {
			clientTimeout = newClientTimeout
		} else {
			fmt.Printf("Error parsing timeout: %s\n", err)
			return
		}
	}
	fmt.Printf("Using a timeout of %d nanoseconds\n", clientTimeout.Nanoseconds())
	if err != nil {
		fmt.Printf("ERROR %s\n", err)
		return
	}
	client := &http.Client{}
	client.Timeout = clientTimeout
	fmt.Printf("Will spawn %d workers making %d requests to %s\n", concurrencycount, maxRequestCount, address)
	runLoadTest(client, sqsurl, address, maxRequestCount, concurrencycount, awsregion)
}

type Job struct{}

type RequestResult struct {
	Time             int64  `json:"time"`
	Host             string `json:"host"`
	Type             string `json:"type"`
	Status           int    `json:"status"`
	ElapsedFirstByte int64  `json:"elapsed-first-byte"`
	ElapsedLastByte  int64  `json:"elapsed-last-byte"`
	Elapsed          int64  `json:"elapsed"`
	Bytes            int    `json:"bytes"`
	Timeout          bool   `json:"timeout"`
	ConnectionError  bool   `json:"connection-error"`
	State            string `json:"state"`
}

func runLoadTest(client *http.Client, sqsurl string, url string, totalRequests int, concurrencycount int, awsregion string) {
	awsConfig := aws.NewConfig().WithRegion(awsregion)
	sqsAdaptor := queue.NewSQSAdaptor(awsConfig, sqsurl)
	//sqsAdaptor := queue.NewDummyAdaptor(sqsurl)
	jobs := make(chan Job, totalRequests)
	ch := make(chan RequestResult, totalRequests)
	var wg sync.WaitGroup
	loadTestStartTime := time.Now()
	var requestsSoFar int
	for i := 0; i < totalRequests; i++ {
		jobs <- Job{}
	}
	close(jobs)
	for i := 0; i < concurrencycount; i++ {
		wg.Add(1)
		go fetch(loadTestStartTime, client, url, totalRequests, jobs, ch, &wg, awsregion)
	}
	fmt.Println("Waiting for results…")

	for requestsSoFar < totalRequests {
		var agg [100]RequestResult
		aggRequestCount := len(agg)
		i := 0

		var timeToFirstTotal int64
		var requestTimeTotal int64
		totBytesRead := 0
		statuses := make(map[string]int)
		var firstRequestTime int64
		var lastRequestTime int64
		var slowest int64
		var fastest int64
		var totalTimedOut int
		var totalConnectionError int
		for ; i < aggRequestCount && requestsSoFar < totalRequests; i++ {
			r := <-ch
			agg[i] = r
			requestsSoFar++
			if requestsSoFar%10 == 0 {
				fmt.Printf("\r%.2f%% done (%d requests out of %d)", (float64(requestsSoFar)/float64(totalRequests))*100.0, requestsSoFar, totalRequests)
			}
			if firstRequestTime == 0 {
				firstRequestTime = r.Time
			}

			lastRequestTime = r.Time

			if r.Timeout {
				totalTimedOut++
				continue
			}
			if r.ConnectionError {
				totalConnectionError++
				continue
			}

			if r.ElapsedLastByte > slowest {
				slowest = r.ElapsedLastByte
			}
			if fastest == 0 {
				fastest = r.ElapsedLastByte
			} else {
				if r.ElapsedLastByte < fastest {
					fastest = r.ElapsedLastByte
				}
			}

			timeToFirstTotal += r.ElapsedFirstByte
			totBytesRead += r.Bytes
			statusStr := strconv.Itoa(r.Status)
			_, ok := statuses[statusStr]
			if !ok {
				statuses[statusStr] = 1
			} else {
				statuses[statusStr]++
			}
			requestTimeTotal += r.Elapsed
		}
		if i == 0 {
			continue
		}
		durationNanoSeconds := lastRequestTime - firstRequestTime
		durationSeconds := float32(durationNanoSeconds) / float32(1000000000)
		aggData := queue.AggData{
			i,
			totalTimedOut,
			totalConnectionError,
			timeToFirstTotal / int64(i),
			totBytesRead,
			statuses,
			requestTimeTotal / int64(i),
			float32(i) / durationSeconds,
			slowest,
			fastest,
			awsregion,
		}
		sqsAdaptor.SendResult(aggData)
	}
	fmt.Printf("\nWaiting for workers…")
	wg.Wait()
	fmt.Printf("\nYay🎈  - %d requests completed\n", requestsSoFar)

}

func fetch(loadTestStartTime time.Time, client *http.Client, address string, requestcount int, jobs <-chan Job, ch chan RequestResult, wg *sync.WaitGroup, awsregion string) {
	defer wg.Done()
	fmt.Printf("Fetching %s\n", address)
	for _ = range jobs {
		start := time.Now()
		req, err := http.NewRequest("GET", address, nil)
		req.Header.Add("User-Agent", "GOAD/0.1")
		response, err := client.Do(req)
		var status string
		var elapsedFirstByte time.Duration
		var elapsedLastByte time.Duration
		var elapsed time.Duration
		var statusCode int
		var bytesRead int
		buf := []byte(" ")
		timedOut := false
		connectionError := false
		if err != nil {
			status = fmt.Sprintf("ERROR: %s\n", err)
			//fmt.Println(status)
			switch err := err.(type) {
			case *url.Error:
				if err, ok := err.Err.(net.Error); ok && err.Timeout() {
					timedOut = true
				}
			case net.Error:
				if err.Timeout() {
					timedOut = true
				}
			}

			if !timedOut {
				connectionError = true
			}
		} else {
			statusCode = response.StatusCode
			_, err = response.Body.Read(buf)
			if err != nil {
				status = fmt.Sprintf("reading first byte failed: %s\n", err)
			}
			elapsedFirstByte = time.Since(start)
			body, err := ioutil.ReadAll(response.Body)
			bytesRead = len(body) + 1
			elapsedLastByte = time.Since(start)
			if err != nil {
				status = fmt.Sprintf("reading response body failed: %s\n", err)
			} else {
				status = "Success"
			}
			elapsed = time.Since(start)
		}
		result := RequestResult{
			start.Sub(loadTestStartTime).Nanoseconds(),
			req.URL.Host,
			req.Method,
			statusCode,
			elapsed.Nanoseconds(),
			elapsedFirstByte.Nanoseconds(),
			elapsedLastByte.Nanoseconds(),
			bytesRead,
			timedOut,
			connectionError,
			status,
		}
		ch <- result
	}
}
<commit_msg>Don't follow redirects, don't treat them as errors<commit_after>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/gophergala2016/goad/queue"
)

func main() {
	address := os.Args[1]
	concurrencycount, err := strconv.Atoi(os.Args[2])
	if err != nil {
		fmt.Printf("ERROR %s\n", err)
		return
	}
	maxRequestCount, err := strconv.Atoi(os.Args[3])
	sqsurl := os.Args[4]
	awsregion := os.Args[5]
	clientTimeout, _ := time.ParseDuration("1s")
	if len(os.Args) > 6 {
		newClientTimeout, err := time.ParseDuration(os.Args[6])
		if err == nil {
			clientTimeout = newClientTimeout
		} else {
			fmt.Printf("Error parsing timeout: %s\n", err)
			return
		}
	}
	fmt.Printf("Using a timeout of %d nanoseconds\n", clientTimeout.Nanoseconds())
	if err != nil {
		fmt.Printf("ERROR %s\n", err)
		return
	}
	client := &http.Client{}
	client.Timeout = clientTimeout
	client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
		return errors.New("redirect")
	}
	fmt.Printf("Will spawn %d workers making %d requests to %s\n", concurrencycount, maxRequestCount, address)
	runLoadTest(client, sqsurl, address, maxRequestCount, concurrencycount, awsregion)
}

type Job struct{}

type RequestResult struct {
	Time             int64  `json:"time"`
	Host             string `json:"host"`
	Type             string `json:"type"`
	Status           int    `json:"status"`
	ElapsedFirstByte int64  `json:"elapsed-first-byte"`
	ElapsedLastByte  int64  `json:"elapsed-last-byte"`
	Elapsed          int64  `json:"elapsed"`
	Bytes            int    `json:"bytes"`
	Timeout          bool   `json:"timeout"`
	ConnectionError  bool   `json:"connection-error"`
	State            string `json:"state"`
}

func runLoadTest(client *http.Client, sqsurl string, url string, totalRequests int, concurrencycount int, awsregion string) {
	awsConfig := aws.NewConfig().WithRegion(awsregion)
	sqsAdaptor := queue.NewSQSAdaptor(awsConfig, sqsurl)
	//sqsAdaptor := queue.NewDummyAdaptor(sqsurl)
	jobs := make(chan Job, totalRequests)
	ch := make(chan RequestResult, totalRequests)
	var wg sync.WaitGroup
	loadTestStartTime := time.Now()
	var requestsSoFar int
	for i := 0; i < totalRequests; i++ {
		jobs <- Job{}
	}
	close(jobs)
	for i := 0; i < concurrencycount; i++ {
		wg.Add(1)
		go fetch(loadTestStartTime, client, url, totalRequests, jobs, ch, &wg, awsregion)
	}
	fmt.Println("Waiting for results…")

	for requestsSoFar < totalRequests {
		var agg [100]RequestResult
		aggRequestCount := len(agg)
		i := 0

		var timeToFirstTotal int64
		var requestTimeTotal int64
		totBytesRead := 0
		statuses := make(map[string]int)
		var firstRequestTime int64
		var lastRequestTime int64
		var slowest int64
		var fastest int64
		var totalTimedOut int
		var totalConnectionError int
		for ; i < aggRequestCount && requestsSoFar < totalRequests; i++ {
			r := <-ch
			agg[i] = r
			requestsSoFar++
			if requestsSoFar%10 == 0 {
				fmt.Printf("\r%.2f%% done (%d requests out of %d)", (float64(requestsSoFar)/float64(totalRequests))*100.0, requestsSoFar, totalRequests)
			}
			if firstRequestTime == 0 {
				firstRequestTime = r.Time
			}

			lastRequestTime = r.Time

			if r.Timeout {
				totalTimedOut++
				continue
			}
			if r.ConnectionError {
				totalConnectionError++
				continue
			}

			if r.ElapsedLastByte > slowest {
				slowest = r.ElapsedLastByte
			}
			if fastest == 0 {
				fastest = r.ElapsedLastByte
			} else {
				if r.ElapsedLastByte < fastest {
					fastest = r.ElapsedLastByte
				}
			}

			timeToFirstTotal += r.ElapsedFirstByte
			totBytesRead += r.Bytes
			statusStr := strconv.Itoa(r.Status)
			_, ok := statuses[statusStr]
			if !ok {
				statuses[statusStr] = 1
			} else {
				statuses[statusStr]++
			}
			requestTimeTotal += r.Elapsed
		}
		if i == 0 {
			continue
		}
		durationNanoSeconds := lastRequestTime - firstRequestTime
		durationSeconds := float32(durationNanoSeconds) / float32(1000000000)
		aggData := queue.AggData{
			i,
			totalTimedOut,
			totalConnectionError,
			timeToFirstTotal / int64(i),
			totBytesRead,
			statuses,
			requestTimeTotal / int64(i),
			float32(i) / durationSeconds,
			slowest,
			fastest,
			awsregion,
		}
		sqsAdaptor.SendResult(aggData)
	}
	fmt.Printf("\nWaiting for workers…")
	wg.Wait()
	fmt.Printf("\nYay🎈  - %d requests completed\n", requestsSoFar)

}

func fetch(loadTestStartTime time.Time, client *http.Client, address string, requestcount int, jobs <-chan Job, ch chan RequestResult, wg *sync.WaitGroup, awsregion string) {
	defer wg.Done()
	fmt.Printf("Fetching %s\n", address)
	for _ = range jobs {
		start := time.Now()
		req, err := http.NewRequest("GET", address, nil)
		req.Header.Add("User-Agent", "GOAD/0.1")
		response, err := client.Do(req)
		var status string
		var elapsedFirstByte time.Duration
		var elapsedLastByte time.Duration
		var elapsed time.Duration
		var statusCode int
		var bytesRead int
		buf := []byte(" ")
		timedOut := false
		connectionError := false
		if err != nil && !strings.Contains(err.Error(), "redirect") {
			status = fmt.Sprintf("ERROR: %s\n", err)
			switch err := err.(type) {
			case *url.Error:
				if err, ok := err.Err.(net.Error); ok && err.Timeout() {
					timedOut = true
				}
			case net.Error:
				if err.Timeout() {
					timedOut = true
				}
			}

			if !timedOut {
				connectionError = true
			}
		} else {
			statusCode = response.StatusCode
			_, err = response.Body.Read(buf)
			if err != nil {
				status = fmt.Sprintf("reading first byte failed: %s\n", err)
			}
			elapsedFirstByte = time.Since(start)
			body, err := ioutil.ReadAll(response.Body)
			bytesRead = len(body) + 1
			elapsedLastByte = time.Since(start)
			if err != nil {
				status = fmt.Sprintf("reading response body failed: %s\n", err)
			} else {
				status = "Success"
			}
			elapsed = time.Since(start)
		}
		result := RequestResult{
			start.Sub(loadTestStartTime).Nanoseconds(),
			req.URL.Host,
			req.Method,
			statusCode,
			elapsed.Nanoseconds(),
			elapsedFirstByte.Nanoseconds(),
			elapsedLastByte.Nanoseconds(),
			bytesRead,
			timedOut,
			connectionError,
			status,
		}
		ch <- result
	}
}
<|endoftext|>


<commit_before>package push

import (
	"crypto/ecdsa"
	"crypto/md5"
	"crypto/tls"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"path/filepath"
	"runtime"
	"time"

	"github.com/cozy/cozy-stack/model/account"
	"github.com/cozy/cozy-stack/model/instance"
	"github.com/cozy/cozy-stack/model/job"
	"github.com/cozy/cozy-stack/model/notification/center"
	"github.com/cozy/cozy-stack/model/oauth"
	"github.com/cozy/cozy-stack/pkg/config/config"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/mail"
	"github.com/sirupsen/logrus"

	fcm "github.com/appleboy/go-fcm"

	apns "github.com/sideshow/apns2"
	apns_cert "github.com/sideshow/apns2/certificate"
	apns_payload "github.com/sideshow/apns2/payload"
	apns_token "github.com/sideshow/apns2/token"
)

var (
	fcmClient *fcm.Client
	iosClient *apns.Client
)

func init() {
	job.AddWorker(&job.WorkerConfig{
		WorkerType:   "push",
		Concurrency:  runtime.NumCPU(),
		MaxExecCount: 1,
		Timeout:      10 * time.Second,
		WorkerInit:   Init,
		WorkerFunc:   Worker,
	})
}

// Init initializes the necessary global clients
func Init() (err error) {
	conf := config.GetConfig().Notifications

	if conf.AndroidAPIKey != "" {
		if conf.FCMServer != "" {
			fcmClient, err = fcm.NewClient(conf.AndroidAPIKey, fcm.WithEndpoint(conf.FCMServer))
		} else {
			fcmClient, err = fcm.NewClient(conf.AndroidAPIKey)
		}
		logger.WithNamespace("push").Infof("Initialized FCM client with Android API Key")
		if err != nil {
			logger.WithNamespace("push").Warnf("%s", err)
			return
		}
	}

	if conf.IOSCertificateKeyPath != "" {
		var authKey *ecdsa.PrivateKey
		var certificateKey tls.Certificate

		switch filepath.Ext(conf.IOSCertificateKeyPath) {
		case ".p12":
			certificateKey, err = apns_cert.FromP12File(
				conf.IOSCertificateKeyPath, conf.IOSCertificatePassword)
		case ".pem":
			certificateKey, err = apns_cert.FromPemFile(
				conf.IOSCertificateKeyPath, conf.IOSCertificatePassword)
		case ".p8":
			authKey, err = apns_token.AuthKeyFromFile(conf.IOSCertificateKeyPath)
		default:
			err = errors.New("wrong certificate key extension")
		}
		if err != nil {
			return err
		}

		if authKey != nil {
			t := &apns_token.Token{
				AuthKey: authKey,
				KeyID:   conf.IOSKeyID,
				TeamID:  conf.IOSTeamID,
			}
			iosClient = apns.NewTokenClient(t)
		} else {
			iosClient = apns.NewClient(certificateKey)
		}
		if conf.Development {
			iosClient = iosClient.Development()
		} else {
			iosClient = iosClient.Production()
		}
	}
	return
}

// Worker is the worker that send push messages.
func Worker(ctx *job.WorkerContext) error {
	var msg center.PushMessage
	if err := ctx.UnmarshalMessage(&msg); err != nil {
		return err
	}
	cs, err := oauth.GetNotifiables(ctx.Instance)
	if err != nil {
		return err
	}
	if len(cs) > 10 {
		ctx.Logger().Warnf("too many notifiable devices: %d", len(cs))
		cs = cs[:10]
	}
	sent := false
	for _, c := range cs {
		if err := push(ctx, c, &msg); err == nil {
			sent = true
		} else {
			ctx.Logger().
				WithFields(logrus.Fields{
					"device_id":       c.ID(),
					"device_platform": c.NotificationPlatform,
				}).
				Warnf("could not send notification on device: %s", err)
		}
	}
	if !sent {
		sendFallbackMail(ctx.Instance, msg.MailFallback)
	}
	return nil
}

func push(ctx *job.WorkerContext, c *oauth.Client, msg *center.PushMessage) error {
	switch c.NotificationPlatform {
	case oauth.PlatformFirebase, "android", "ios":
		return pushToFirebase(ctx, c, msg)
	case oauth.PlatformAPNS:
		return pushToAPNS(ctx, c, msg)
	default:
		return fmt.Errorf("notifications: unknown platform %q", c.NotificationPlatform)
	}
}

// Firebase Cloud Messaging HTTP Protocol
// https://firebase.google.com/docs/cloud-messaging/http-server-ref
func pushToFirebase(ctx *job.WorkerContext, c *oauth.Client, msg *center.PushMessage) error {
	client := getFirebaseClient(msg.Slug(), ctx.Instance.ContextName)

	if client == nil {
		ctx.Logger().Warn("Could not send android notification: not configured")
		return nil
	}

	var priority string
	if msg.Priority == "high" {
		priority = "high"
	}

	var hashedSource []byte
	if msg.Collapsible {
		hashedSource = hashSource(msg.Source)
	} else {
		hashedSource = hashSource(msg.Source + msg.NotificationID)
	}

	// notID should be an integer, we take the first 32bits of the hashed source
	// value.
	notID := int32(binary.BigEndian.Uint32(hashedSource[:4]))
	if notID < 0 {
		notID = -notID
	}

	notification := &fcm.Message{
		To:               c.NotificationDeviceToken,
		Priority:         priority,
		ContentAvailable: true,
		Notification: &fcm.Notification{
			Sound: msg.Sound,
			Title: msg.Title,
			Body:  msg.Message,
		},
		Data: map[string]interface{}{
			// Fields required by phonegap-plugin-push
			// see: https://github.com/phonegap/phonegap-plugin-push/blob/master/docs/PAYLOAD.md#android-behaviour
			"notId": notID,
			"title": msg.Title,
			"body":  msg.Message,
		},
	}

	ctx.Logger().Infof("Built notification for FCM: %#v", notification)
	if msg.Collapsible {
		notification.CollapseKey = hex.EncodeToString(hashedSource)
	}
	for k, v := range msg.Data {
		notification.Data[k] = v
	}

	ctx.Logger().Infof("FCM send: %#v", notification)
	res, err := client.Send(notification)
	if err != nil {
		ctx.Logger().Warnf("Error during fcm send: %s", err)
		return err
	}
	if res.Failure == 0 {
		return nil
	}

	for _, result := range res.Results {
		if err = result.Error; err != nil {
			return err
		}
	}
	return nil
}

func getFirebaseClient(slug, contextName string) *fcm.Client {
	if slug == "" {
		return fcmClient
	}
	typ, err := account.TypeInfo(slug, contextName)
	if err == nil && typ.AndroidAPIKey != "" {
		client, err := fcm.NewClient(typ.AndroidAPIKey)
		if err != nil {
			return nil
		}
		return client
	}
	return fcmClient
}

func pushToAPNS(ctx *job.WorkerContext, c *oauth.Client, msg *center.PushMessage) error {
	if iosClient == nil {
		ctx.Logger().Warn("Could not send iOS notification: not configured")
		return nil
	}

	var priority int
	if msg.Priority == "normal" {
		priority = apns.PriorityLow
	} else {
		priority = apns.PriorityHigh
	}

	payload := apns_payload.NewPayload().
		AlertTitle(msg.Title).
		Alert(msg.Message).
		Sound(msg.Sound)

	for k, v := range msg.Data {
		payload.Custom(k, v)
	}

	notification := &apns.Notification{
		DeviceToken: c.NotificationDeviceToken,
		Payload:     payload,
		Priority:    priority,
		CollapseID:  hex.EncodeToString(hashSource(msg.Source)), // CollapseID should not exceed 64 bytes
	}

	res, err := iosClient.PushWithContext(ctx, notification)
	if err != nil {
		return err
	}
	if res.StatusCode != 200 {
		return fmt.Errorf("failed to push apns notification: %d %s", res.StatusCode, res.Reason)
	}
	return nil
}

func hashSource(source string) []byte {
	h := md5.New()
	_, _ = h.Write([]byte(source))
	return h.Sum(nil)
}

func sendFallbackMail(inst *instance.Instance, email *mail.Options) {
	if inst == nil || email == nil {
		return
	}
	msg, err := job.NewMessage(&email)
	if err != nil {
		return
	}
	_, _ = job.System().PushJob(inst, &job.JobRequest{
		WorkerType: "sendmail",
		Message:    msg,
	})
}
<commit_msg>Avoid duplicate notification (#2929)<commit_after>package push

import (
	"crypto/ecdsa"
	"crypto/md5"
	"crypto/tls"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"path/filepath"
	"runtime"
	"time"

	"github.com/cozy/cozy-stack/model/account"
	"github.com/cozy/cozy-stack/model/instance"
	"github.com/cozy/cozy-stack/model/job"
	"github.com/cozy/cozy-stack/model/notification/center"
	"github.com/cozy/cozy-stack/model/oauth"
	"github.com/cozy/cozy-stack/pkg/config/config"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/mail"
	"github.com/sirupsen/logrus"

	fcm "github.com/appleboy/go-fcm"

	apns "github.com/sideshow/apns2"
	apns_cert "github.com/sideshow/apns2/certificate"
	apns_payload "github.com/sideshow/apns2/payload"
	apns_token "github.com/sideshow/apns2/token"
)

var (
	fcmClient *fcm.Client
	iosClient *apns.Client
)

func init() {
	job.AddWorker(&job.WorkerConfig{
		WorkerType:   "push",
		Concurrency:  runtime.NumCPU(),
		MaxExecCount: 1,
		Timeout:      10 * time.Second,
		WorkerInit:   Init,
		WorkerFunc:   Worker,
	})
}

// Init initializes the necessary global clients
func Init() (err error) {
	conf := config.GetConfig().Notifications

	if conf.AndroidAPIKey != "" {
		if conf.FCMServer != "" {
			fcmClient, err = fcm.NewClient(conf.AndroidAPIKey, fcm.WithEndpoint(conf.FCMServer))
		} else {
			fcmClient, err = fcm.NewClient(conf.AndroidAPIKey)
		}
		logger.WithNamespace("push").Infof("Initialized FCM client with Android API Key")
		if err != nil {
			logger.WithNamespace("push").Warnf("%s", err)
			return
		}
	}

	if conf.IOSCertificateKeyPath != "" {
		var authKey *ecdsa.PrivateKey
		var certificateKey tls.Certificate

		switch filepath.Ext(conf.IOSCertificateKeyPath) {
		case ".p12":
			certificateKey, err = apns_cert.FromP12File(
				conf.IOSCertificateKeyPath, conf.IOSCertificatePassword)
		case ".pem":
			certificateKey, err = apns_cert.FromPemFile(
				conf.IOSCertificateKeyPath, conf.IOSCertificatePassword)
		case ".p8":
			authKey, err = apns_token.AuthKeyFromFile(conf.IOSCertificateKeyPath)
		default:
			err = errors.New("wrong certificate key extension")
		}
		if err != nil {
			return err
		}

		if authKey != nil {
			t := &apns_token.Token{
				AuthKey: authKey,
				KeyID:   conf.IOSKeyID,
				TeamID:  conf.IOSTeamID,
			}
			iosClient = apns.NewTokenClient(t)
		} else {
			iosClient = apns.NewClient(certificateKey)
		}
		if conf.Development {
			iosClient = iosClient.Development()
		} else {
			iosClient = iosClient.Production()
		}
	}
	return
}

// Worker is the worker that send push messages.
func Worker(ctx *job.WorkerContext) error {
	var msg center.PushMessage
	if err := ctx.UnmarshalMessage(&msg); err != nil {
		return err
	}
	cs, err := oauth.GetNotifiables(ctx.Instance)
	if err != nil {
		return err
	}
	if len(cs) > 10 {
		ctx.Logger().Warnf("too many notifiable devices: %d", len(cs))
		cs = cs[:10]
	}
	sent := false
	seen := make(map[string]struct{})
	for _, c := range cs {
		if _, ok := seen[c.NotificationDeviceToken]; ok {
			continue
		}
		seen[c.NotificationDeviceToken] = struct{}{}
		if err := push(ctx, c, &msg); err == nil {
			sent = true
		} else {
			ctx.Logger().
				WithFields(logrus.Fields{
					"device_id":       c.ID(),
					"device_platform": c.NotificationPlatform,
				}).
				Warnf("could not send notification on device: %s", err)
		}
	}
	if !sent {
		sendFallbackMail(ctx.Instance, msg.MailFallback)
	}
	return nil
}

func push(ctx *job.WorkerContext, c *oauth.Client, msg *center.PushMessage) error {
	switch c.NotificationPlatform {
	case oauth.PlatformFirebase, "android", "ios":
		return pushToFirebase(ctx, c, msg)
	case oauth.PlatformAPNS:
		return pushToAPNS(ctx, c, msg)
	default:
		return fmt.Errorf("notifications: unknown platform %q", c.NotificationPlatform)
	}
}

// Firebase Cloud Messaging HTTP Protocol
// https://firebase.google.com/docs/cloud-messaging/http-server-ref
func pushToFirebase(ctx *job.WorkerContext, c *oauth.Client, msg *center.PushMessage) error {
	client := getFirebaseClient(msg.Slug(), ctx.Instance.ContextName)

	if client == nil {
		ctx.Logger().Warn("Could not send android notification: not configured")
		return nil
	}

	var priority string
	if msg.Priority == "high" {
		priority = "high"
	}

	var hashedSource []byte
	if msg.Collapsible {
		hashedSource = hashSource(msg.Source)
	} else {
		hashedSource = hashSource(msg.Source + msg.NotificationID)
	}

	// notID should be an integer, we take the first 32bits of the hashed source
	// value.
	notID := int32(binary.BigEndian.Uint32(hashedSource[:4]))
	if notID < 0 {
		notID = -notID
	}

	notification := &fcm.Message{
		To:               c.NotificationDeviceToken,
		Priority:         priority,
		ContentAvailable: true,
		Notification: &fcm.Notification{
			Sound: msg.Sound,
			Title: msg.Title,
			Body:  msg.Message,
		},
		Data: map[string]interface{}{
			// Fields required by phonegap-plugin-push
			// see: https://github.com/phonegap/phonegap-plugin-push/blob/master/docs/PAYLOAD.md#android-behaviour
			"notId": notID,
			"title": msg.Title,
			"body":  msg.Message,
		},
	}

	ctx.Logger().Infof("Built notification for FCM: %#v", notification)
	if msg.Collapsible {
		notification.CollapseKey = hex.EncodeToString(hashedSource)
	}
	for k, v := range msg.Data {
		notification.Data[k] = v
	}

	ctx.Logger().Infof("FCM send: %#v", notification)
	res, err := client.Send(notification)
	if err != nil {
		ctx.Logger().Warnf("Error during fcm send: %s", err)
		return err
	}
	if res.Failure == 0 {
		return nil
	}

	for _, result := range res.Results {
		if err = result.Error; err != nil {
			return err
		}
	}
	return nil
}

func getFirebaseClient(slug, contextName string) *fcm.Client {
	if slug == "" {
		return fcmClient
	}
	typ, err := account.TypeInfo(slug, contextName)
	if err == nil && typ.AndroidAPIKey != "" {
		client, err := fcm.NewClient(typ.AndroidAPIKey)
		if err != nil {
			return nil
		}
		return client
	}
	return fcmClient
}

func pushToAPNS(ctx *job.WorkerContext, c *oauth.Client, msg *center.PushMessage) error {
	if iosClient == nil {
		ctx.Logger().Warn("Could not send iOS notification: not configured")
		return nil
	}

	var priority int
	if msg.Priority == "normal" {
		priority = apns.PriorityLow
	} else {
		priority = apns.PriorityHigh
	}

	payload := apns_payload.NewPayload().
		AlertTitle(msg.Title).
		Alert(msg.Message).
		Sound(msg.Sound)

	for k, v := range msg.Data {
		payload.Custom(k, v)
	}

	notification := &apns.Notification{
		DeviceToken: c.NotificationDeviceToken,
		Payload:     payload,
		Priority:    priority,
		CollapseID:  hex.EncodeToString(hashSource(msg.Source)), // CollapseID should not exceed 64 bytes
	}

	res, err := iosClient.PushWithContext(ctx, notification)
	if err != nil {
		return err
	}
	if res.StatusCode != 200 {
		return fmt.Errorf("failed to push apns notification: %d %s", res.StatusCode, res.Reason)
	}
	return nil
}

func hashSource(source string) []byte {
	h := md5.New()
	_, _ = h.Write([]byte(source))
	return h.Sum(nil)
}

func sendFallbackMail(inst *instance.Instance, email *mail.Options) {
	if inst == nil || email == nil {
		return
	}
	msg, err := job.NewMessage(&email)
	if err != nil {
		return
	}
	_, _ = job.System().PushJob(inst, &job.JobRequest{
		WorkerType: "sendmail",
		Message:    msg,
	})
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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 fitasks

import (
	"crypto/x509"
	"fmt"
	"net"
	"sort"
	"strings"

	"github.com/golang/glog"
	"k8s.io/kops/pkg/pki"
	"k8s.io/kops/upup/pkg/fi"
)

var wellKnownCertificateTypes = map[string]string{
	"ca":           "CA,KeyUsageCRLSign,KeyUsageCertSign",
	"client":       "ExtKeyUsageClientAuth,KeyUsageDigitalSignature",
	"clientServer": "ExtKeyUsageServerAuth,KeyUsageDigitalSignature,ExtKeyUsageClientAuth,KeyUsageKeyEncipherment",
	"server":       "ExtKeyUsageServerAuth,KeyUsageDigitalSignature,KeyUsageKeyEncipherment",
}

//go:generate fitask -type=Keypair
type Keypair struct {
	// Name is the name of the keypair
	Name *string
	// AlternateNames a list of alternative names for this certificate
	AlternateNames []string `json:"alternateNames"`
	// AlternateNameTasks is a collection of subtask
	AlternateNameTasks []fi.Task `json:"alternateNameTasks"`
	// Lifecycle is context for a task
	Lifecycle *fi.Lifecycle
	// Signer is the keypair to use to sign, for when we want to use an alternative CA
	Signer *Keypair
	// Subject is the cerificate subject
	Subject string `json:"subject"`
	// Type the type of certificate i.e. CA, server, client etc
	Type string `json:"type"`
	// Format stores the api version of kops.Keyset.  We are using this info in order to determine if kops
	// is accessing legacy secrets that do not use keyset.yaml.
	Format string `json:"format"`
}

var _ fi.HasCheckExisting = &Keypair{}
var _ fi.HasName = &Keypair{}

// It's important always to check for the existing key, so we don't regenerate keys e.g. on terraform
func (e *Keypair) CheckExisting(c *fi.Context) bool {
	return true
}

var _ fi.CompareWithID = &Keypair{}

func (e *Keypair) CompareWithID() *string {
	return &e.Subject
}

func (e *Keypair) Find(c *fi.Context) (*Keypair, error) {
	name := fi.StringValue(e.Name)
	if name == "" {
		return nil, nil
	}

	cert, key, format, err := c.Keystore.FindKeypair(name)
	if err != nil {
		return nil, err
	}
	if cert == nil {
		return nil, nil
	}
	if key == nil {
		return nil, fmt.Errorf("found cert in store, but did not find private key: %q", name)
	}

	var alternateNames []string
	alternateNames = append(alternateNames, cert.Certificate.DNSNames...)
	alternateNames = append(alternateNames, cert.Certificate.EmailAddresses...)
	for _, ip := range cert.Certificate.IPAddresses {
		alternateNames = append(alternateNames, ip.String())
	}
	sort.Strings(alternateNames)

	actual := &Keypair{
		Name:           &name,
		AlternateNames: alternateNames,
		Subject:        pkixNameToString(&cert.Subject),
		Type:           buildTypeDescription(cert.Certificate),
		Format:         string(format),
	}

	actual.Signer = &Keypair{Subject: pkixNameToString(&cert.Certificate.Issuer)}

	// Avoid spurious changes
	actual.Lifecycle = e.Lifecycle

	return actual, nil
}

func (e *Keypair) Run(c *fi.Context) error {
	err := e.normalize(c)
	if err != nil {
		return err
	}
	return fi.DefaultDeltaRunMethod(e, c)
}

func (e *Keypair) normalize(c *fi.Context) error {
	var alternateNames []string

	for _, s := range e.AlternateNames {
		s = strings.TrimSpace(s)
		if s == "" {
			continue
		}
		alternateNames = append(alternateNames, s)
	}

	for _, task := range e.AlternateNameTasks {
		if hasAddress, ok := task.(fi.HasAddress); ok {
			address, err := hasAddress.FindIPAddress(c)
			if err != nil {
				return fmt.Errorf("error finding address for %v: %v", task, err)
			}
			if address == nil {
				glog.Warningf("Task did not have an address: %v", task)
				continue
			}
			glog.V(8).Infof("Resolved alternateName %q for %q", *address, task)
			alternateNames = append(alternateNames, *address)
		} else {
			return fmt.Errorf("Unsupported type for AlternateNameDependencies: %v", task)
		}
	}

	sort.Strings(alternateNames)
	e.AlternateNames = alternateNames
	e.AlternateNameTasks = nil

	return nil
}

func (_ *Keypair) CheckChanges(a, e, changes *Keypair) error {
	if a != nil {
		if changes.Name != nil {
			return fi.CannotChangeField("Name")
		}
	}
	return nil
}

func (_ *Keypair) Render(c *fi.Context, a, e, changes *Keypair) error {
	name := fi.StringValue(e.Name)
	if name == "" {
		return fi.RequiredField("Name")
	}

	template, err := e.BuildCertificateTemplate()
	if err != nil {
		return err
	}

	changeStoredFormat := false
	createCertificate := false
	if a == nil {
		createCertificate = true
		glog.V(8).Infof("creating brand new certificate")
	} else if changes != nil {
		glog.V(8).Infof("creating certificate as changes are not nil")
		if changes.AlternateNames != nil {
			createCertificate = true
			glog.V(8).Infof("creating certificate new AlternateNames")
		} else if changes.Subject != "" {
			createCertificate = true
			glog.V(8).Infof("creating certificate new Subject")
		} else if changes.Type != "" {
			createCertificate = true
			glog.V(8).Infof("creating certificate new Type")
		} else if changes.Format != "" {
			changeStoredFormat = true
		} else {
			glog.Warningf("Ignoring changes in key: %v", fi.DebugAsJsonString(changes))
		}
	}

	if createCertificate {
		glog.V(2).Infof("Creating PKI keypair %q", name)

		cert, privateKey, _, err := c.Keystore.FindKeypair(name)
		if err != nil {
			return err
		}

		// We always reuse the private key if it exists,
		// if we change keys we often have to regenerate e.g. the service accounts
		// TODO: Eventually rotate keys / don't always reuse?
		if privateKey == nil {
			glog.V(2).Infof("Creating privateKey %q", name)

			privateKey, err = pki.GeneratePrivateKey()
			if err != nil {
				return err
			}
		}

		signer := fi.CertificateId_CA
		if e.Signer != nil {
			signer = fi.StringValue(e.Signer.Name)
		}

		cert, err = c.Keystore.CreateKeypair(signer, name, template, privateKey)
		if err != nil {
			return err
		}

		glog.V(8).Infof("created certificate with cn=%s", cert.Subject.CommonName)
	}

	// TODO: Check correct subject / flags

	if changeStoredFormat {
		// We fetch and reinsert the same keypair, forcing an update to our preferred format
		// TODO: We're assuming that we want to save in the preferred format
		cert, privateKey, _, err := c.Keystore.FindKeypair(name)
		if err != nil {
			return err
		}
		err = c.Keystore.StoreKeypair(name, cert, privateKey)
		if err != nil {
			return err
		}

		glog.Infof("updated Keypair %q to API format %q", name, e.Format)
	}

	return nil
}

// BuildCertificateTemplate is responsible for constructing a certificate template
func (e *Keypair) BuildCertificateTemplate() (*x509.Certificate, error) {
	template, err := buildCertificateTemplateForType(e.Type)
	if err != nil {
		return nil, err
	}

	subjectPkix, err := parsePkixName(e.Subject)
	if err != nil {
		return nil, fmt.Errorf("error parsing Subject: %v", err)
	}

	if len(subjectPkix.ToRDNSequence()) == 0 {
		return nil, fmt.Errorf("Subject name was empty for SSL keypair %q", e.Name)
	}

	template.Subject = *subjectPkix

	var alternateNames []string
	alternateNames = append(alternateNames, e.AlternateNames...)

	for _, san := range alternateNames {
		san = strings.TrimSpace(san)
		if san == "" {
			continue
		}
		if ip := net.ParseIP(san); ip != nil {
			template.IPAddresses = append(template.IPAddresses, ip)
		} else {
			template.DNSNames = append(template.DNSNames, san)
		}
	}

	return template, nil
}

func buildCertificateTemplateForType(certificateType string) (*x509.Certificate, error) {
	if expanded, found := wellKnownCertificateTypes[certificateType]; found {
		certificateType = expanded
	}

	template := &x509.Certificate{
		BasicConstraintsValid: true,
		IsCA: false,
	}

	tokens := strings.Split(certificateType, ",")
	for _, t := range tokens {
		if strings.HasPrefix(t, "KeyUsage") {
			ku, found := parseKeyUsage(t)
			if !found {
				return nil, fmt.Errorf("unrecognized certificate option: %v", t)
			}
			template.KeyUsage |= ku
		} else if strings.HasPrefix(t, "ExtKeyUsage") {
			ku, found := parseExtKeyUsage(t)
			if !found {
				return nil, fmt.Errorf("unrecognized certificate option: %v", t)
			}
			template.ExtKeyUsage = append(template.ExtKeyUsage, ku)
		} else if t == "CA" {
			template.IsCA = true
		} else {
			return nil, fmt.Errorf("unrecognized certificate option: %q", t)
		}
	}

	return template, nil
}

// buildTypeDescription extracts the type based on the certificate extensions
func buildTypeDescription(cert *x509.Certificate) string {
	var options []string

	if cert.IsCA {
		options = append(options, "CA")
	}

	options = append(options, keyUsageToString(cert.KeyUsage)...)

	for _, extKeyUsage := range cert.ExtKeyUsage {
		options = append(options, extKeyUsageToString(extKeyUsage))
	}

	sort.Strings(options)
	s := strings.Join(options, ",")

	for k, v := range wellKnownCertificateTypes {
		if v == s {
			s = k
		}
	}

	return s
}
<commit_msg>Fix etcd Keypair change showing on every kops update when TLS enabled<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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 fitasks

import (
	"crypto/x509"
	"fmt"
	"net"
	"sort"
	"strings"

	"github.com/golang/glog"
	"k8s.io/kops/pkg/pki"
	"k8s.io/kops/upup/pkg/fi"
)

var wellKnownCertificateTypes = map[string]string{
	"ca":           "CA,KeyUsageCRLSign,KeyUsageCertSign",
	"client":       "ExtKeyUsageClientAuth,KeyUsageDigitalSignature",
	"clientServer": "ExtKeyUsageClientAuth,ExtKeyUsageServerAuth,KeyUsageDigitalSignature,KeyUsageKeyEncipherment",
	"server":       "ExtKeyUsageServerAuth,KeyUsageDigitalSignature,KeyUsageKeyEncipherment",
}

//go:generate fitask -type=Keypair
type Keypair struct {
	// Name is the name of the keypair
	Name *string
	// AlternateNames a list of alternative names for this certificate
	AlternateNames []string `json:"alternateNames"`
	// AlternateNameTasks is a collection of subtask
	AlternateNameTasks []fi.Task `json:"alternateNameTasks"`
	// Lifecycle is context for a task
	Lifecycle *fi.Lifecycle
	// Signer is the keypair to use to sign, for when we want to use an alternative CA
	Signer *Keypair
	// Subject is the cerificate subject
	Subject string `json:"subject"`
	// Type the type of certificate i.e. CA, server, client etc
	Type string `json:"type"`
	// Format stores the api version of kops.Keyset.  We are using this info in order to determine if kops
	// is accessing legacy secrets that do not use keyset.yaml.
	Format string `json:"format"`
}

var _ fi.HasCheckExisting = &Keypair{}
var _ fi.HasName = &Keypair{}

// It's important always to check for the existing key, so we don't regenerate keys e.g. on terraform
func (e *Keypair) CheckExisting(c *fi.Context) bool {
	return true
}

var _ fi.CompareWithID = &Keypair{}

func (e *Keypair) CompareWithID() *string {
	return &e.Subject
}

func (e *Keypair) Find(c *fi.Context) (*Keypair, error) {
	name := fi.StringValue(e.Name)
	if name == "" {
		return nil, nil
	}

	cert, key, format, err := c.Keystore.FindKeypair(name)
	if err != nil {
		return nil, err
	}
	if cert == nil {
		return nil, nil
	}
	if key == nil {
		return nil, fmt.Errorf("found cert in store, but did not find private key: %q", name)
	}

	var alternateNames []string
	alternateNames = append(alternateNames, cert.Certificate.DNSNames...)
	alternateNames = append(alternateNames, cert.Certificate.EmailAddresses...)
	for _, ip := range cert.Certificate.IPAddresses {
		alternateNames = append(alternateNames, ip.String())
	}
	sort.Strings(alternateNames)

	actual := &Keypair{
		Name:           &name,
		AlternateNames: alternateNames,
		Subject:        pkixNameToString(&cert.Subject),
		Type:           buildTypeDescription(cert.Certificate),
		Format:         string(format),
	}

	actual.Signer = &Keypair{Subject: pkixNameToString(&cert.Certificate.Issuer)}

	// Avoid spurious changes
	actual.Lifecycle = e.Lifecycle

	return actual, nil
}

func (e *Keypair) Run(c *fi.Context) error {
	err := e.normalize(c)
	if err != nil {
		return err
	}
	return fi.DefaultDeltaRunMethod(e, c)
}

func (e *Keypair) normalize(c *fi.Context) error {
	var alternateNames []string

	for _, s := range e.AlternateNames {
		s = strings.TrimSpace(s)
		if s == "" {
			continue
		}
		alternateNames = append(alternateNames, s)
	}

	for _, task := range e.AlternateNameTasks {
		if hasAddress, ok := task.(fi.HasAddress); ok {
			address, err := hasAddress.FindIPAddress(c)
			if err != nil {
				return fmt.Errorf("error finding address for %v: %v", task, err)
			}
			if address == nil {
				glog.Warningf("Task did not have an address: %v", task)
				continue
			}
			glog.V(8).Infof("Resolved alternateName %q for %q", *address, task)
			alternateNames = append(alternateNames, *address)
		} else {
			return fmt.Errorf("Unsupported type for AlternateNameDependencies: %v", task)
		}
	}

	sort.Strings(alternateNames)
	e.AlternateNames = alternateNames
	e.AlternateNameTasks = nil

	return nil
}

func (_ *Keypair) CheckChanges(a, e, changes *Keypair) error {
	if a != nil {
		if changes.Name != nil {
			return fi.CannotChangeField("Name")
		}
	}
	return nil
}

func (_ *Keypair) Render(c *fi.Context, a, e, changes *Keypair) error {
	name := fi.StringValue(e.Name)
	if name == "" {
		return fi.RequiredField("Name")
	}

	template, err := e.BuildCertificateTemplate()
	if err != nil {
		return err
	}

	changeStoredFormat := false
	createCertificate := false
	if a == nil {
		createCertificate = true
		glog.V(8).Infof("creating brand new certificate")
	} else if changes != nil {
		glog.V(8).Infof("creating certificate as changes are not nil")
		if changes.AlternateNames != nil {
			createCertificate = true
			glog.V(8).Infof("creating certificate new AlternateNames")
		} else if changes.Subject != "" {
			createCertificate = true
			glog.V(8).Infof("creating certificate new Subject")
		} else if changes.Type != "" {
			createCertificate = true
			glog.V(8).Infof("creating certificate new Type")
		} else if changes.Format != "" {
			changeStoredFormat = true
		} else {
			glog.Warningf("Ignoring changes in key: %v", fi.DebugAsJsonString(changes))
		}
	}

	if createCertificate {
		glog.V(2).Infof("Creating PKI keypair %q", name)

		cert, privateKey, _, err := c.Keystore.FindKeypair(name)
		if err != nil {
			return err
		}

		// We always reuse the private key if it exists,
		// if we change keys we often have to regenerate e.g. the service accounts
		// TODO: Eventually rotate keys / don't always reuse?
		if privateKey == nil {
			glog.V(2).Infof("Creating privateKey %q", name)

			privateKey, err = pki.GeneratePrivateKey()
			if err != nil {
				return err
			}
		}

		signer := fi.CertificateId_CA
		if e.Signer != nil {
			signer = fi.StringValue(e.Signer.Name)
		}

		cert, err = c.Keystore.CreateKeypair(signer, name, template, privateKey)
		if err != nil {
			return err
		}

		glog.V(8).Infof("created certificate with cn=%s", cert.Subject.CommonName)
	}

	// TODO: Check correct subject / flags

	if changeStoredFormat {
		// We fetch and reinsert the same keypair, forcing an update to our preferred format
		// TODO: We're assuming that we want to save in the preferred format
		cert, privateKey, _, err := c.Keystore.FindKeypair(name)
		if err != nil {
			return err
		}
		err = c.Keystore.StoreKeypair(name, cert, privateKey)
		if err != nil {
			return err
		}

		glog.Infof("updated Keypair %q to API format %q", name, e.Format)
	}

	return nil
}

// BuildCertificateTemplate is responsible for constructing a certificate template
func (e *Keypair) BuildCertificateTemplate() (*x509.Certificate, error) {
	template, err := buildCertificateTemplateForType(e.Type)
	if err != nil {
		return nil, err
	}

	subjectPkix, err := parsePkixName(e.Subject)
	if err != nil {
		return nil, fmt.Errorf("error parsing Subject: %v", err)
	}

	if len(subjectPkix.ToRDNSequence()) == 0 {
		return nil, fmt.Errorf("Subject name was empty for SSL keypair %q", e.Name)
	}

	template.Subject = *subjectPkix

	var alternateNames []string
	alternateNames = append(alternateNames, e.AlternateNames...)

	for _, san := range alternateNames {
		san = strings.TrimSpace(san)
		if san == "" {
			continue
		}
		if ip := net.ParseIP(san); ip != nil {
			template.IPAddresses = append(template.IPAddresses, ip)
		} else {
			template.DNSNames = append(template.DNSNames, san)
		}
	}

	return template, nil
}

func buildCertificateTemplateForType(certificateType string) (*x509.Certificate, error) {
	if expanded, found := wellKnownCertificateTypes[certificateType]; found {
		certificateType = expanded
	}

	template := &x509.Certificate{
		BasicConstraintsValid: true,
		IsCA: false,
	}

	tokens := strings.Split(certificateType, ",")
	for _, t := range tokens {
		if strings.HasPrefix(t, "KeyUsage") {
			ku, found := parseKeyUsage(t)
			if !found {
				return nil, fmt.Errorf("unrecognized certificate option: %v", t)
			}
			template.KeyUsage |= ku
		} else if strings.HasPrefix(t, "ExtKeyUsage") {
			ku, found := parseExtKeyUsage(t)
			if !found {
				return nil, fmt.Errorf("unrecognized certificate option: %v", t)
			}
			template.ExtKeyUsage = append(template.ExtKeyUsage, ku)
		} else if t == "CA" {
			template.IsCA = true
		} else {
			return nil, fmt.Errorf("unrecognized certificate option: %q", t)
		}
	}

	return template, nil
}

// buildTypeDescription extracts the type based on the certificate extensions
func buildTypeDescription(cert *x509.Certificate) string {
	var options []string

	if cert.IsCA {
		options = append(options, "CA")
	}

	options = append(options, keyUsageToString(cert.KeyUsage)...)

	for _, extKeyUsage := range cert.ExtKeyUsage {
		options = append(options, extKeyUsageToString(extKeyUsage))
	}

	sort.Strings(options)
	s := strings.Join(options, ",")

	for k, v := range wellKnownCertificateTypes {
		if v == s {
			s = k
		}
	}

	return s
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net"
	"strings"

	pb "github.com/shijuvar/go-recipes/grpc/customer"

	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

const (
	port = ":50051"
)

// server is used to implement customer.CustomerServer.
type server struct {
	savedCustomers []*pb.CustomerRequest
}

// CreateCustomer creates a new Customer
func (s *server) CreateCustomer(ctx context.Context, in *pb.CustomerRequest) (*pb.CustomerResponse, error) {
	s.savedCustomers = append(s.savedCustomers, in)
	return &pb.CustomerResponse{Id: in.Id, Success: true}, nil
}

// GetCustomers returns all customers by given filter
func (s *server) GetCustomers(filter *pb.CustomerFilter, stream pb.Customer_GetCustomersServer) error {
	for _, customer := range s.savedCustomers {
		if filter.Keyword != "" {
			if !strings.Contains(customer.Name, filter.Keyword) {
				continue
			}
		}
		if err := stream.Send(customer); err != nil {
			return err
		}
	}
	return nil
}

func main() {
	lis, err := net.Listen("tcp", port)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	// Creates a new gRPC server
	s := grpc.NewServer()
	pb.RegisterCustomerServer(s, &server{})
	s.Serve(lis)
}
<commit_msg>Minor modification<commit_after>package main

import (
	"log"
	"net"
	"strings"

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

	pb "github.com/shijuvar/go-recipes/grpc/customer"
)

const (
	port = ":50051"
)

// server is used to implement customer.CustomerServer.
type server struct {
	savedCustomers []*pb.CustomerRequest
}

// CreateCustomer creates a new Customer
func (s *server) CreateCustomer(ctx context.Context, in *pb.CustomerRequest) (*pb.CustomerResponse, error) {
	s.savedCustomers = append(s.savedCustomers, in)
	return &pb.CustomerResponse{Id: in.Id, Success: true}, nil
}

// GetCustomers returns all customers by given filter
func (s *server) GetCustomers(filter *pb.CustomerFilter, stream pb.Customer_GetCustomersServer) error {
	for _, customer := range s.savedCustomers {
		if filter.Keyword != "" {
			if !strings.Contains(customer.Name, filter.Keyword) {
				continue
			}
		}
		if err := stream.Send(customer); err != nil {
			return err
		}
	}
	return nil
}

func main() {
	lis, err := net.Listen("tcp", port)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	// Creates a new gRPC server
	s := grpc.NewServer()
	pb.RegisterCustomerServer(s, &server{})
	s.Serve(lis)
}
<|endoftext|>


<commit_before>package workers

import (
	"strings"
	"time"

	ogoutils "github.com/Odinman/ogo/utils"
	"github.com/Odinman/omq/utils"
	zmq "github.com/pebbe/zmq4"
)

/* {{{ func connectQueue() (*zmq.Socket, *zmq.Poller)
 *  Helper function that returns a new configured socket
 *  connected to the Paranoid Pirate queue
 */
func (w *OmqWorker) connectQueue() (*zmq.Socket, *zmq.Poller) {
	soc, _ := zmq.NewSocket(zmq.DEALER)
	soc.Connect("inproc://backend")

	//  Tell queue we're ready for work
	soc.Send(PPP_READY, 0)

	poller := zmq.NewPoller()
	poller.Add(soc, zmq.POLLIN)

	return soc, poller
}

/* }}} */

/* {{{ func (w *OmqWorker) newResponser(i int)
 * 回复节点
 */
func (w *OmqWorker) newResponser(i int) {
	node, poller := w.connectQueue()
	w.Trace("%d node ready", i)

	//  If liveness hits zero, queue is considered disconnected
	liveness := HEARTBEAT_LIVENESS
	interval := INTERVAL_INIT

	//  Send out heartbeats at regular intervals
	heartbeat_at := time.Tick(HEARTBEAT_INTERVAL)

	lastCycles := 0
	for cycles := 0; true; {
		sockets, err := poller.Poll(HEARTBEAT_INTERVAL)
		if err != nil {
			w.Error("polling error: %s", err)
			break //  Interrupted
		}

		if len(sockets) == 1 {
			//  Get message
			//  - n-part envelope + content -> request
			//  - 1-part HEARTBEAT -> heartbeat
			msg, err := node.RecvMessage(0)
			if err != nil {
				w.Error("recv error: %s", err)
				break //  Interrupted
			}

			client, cmd := utils.Unwrap(msg)
			if len(msg) >= 4 { //命令应该大于5帧(包含信封以及空帧)
				cycles++

				w.Trace("recv cmd: %s, from client: %q", cmd, client)

				act := strings.ToUpper(cmd[0])
				key := cmd[1]
				switch act {
				case COMMAND_GET, COMMAND_TIMING: //获取key内容
					if r, err := w.localGet(cmd); err != nil {
						w.Debug("error: %s", err)
						if err == ErrNil {
							node.SendMessage(client, "", RESPONSE_NIL) //没有内容,返回空
						} else {
							node.SendMessage(client, "", RESPONSE_ERROR, err.Error()) //回复REQ,因此要加上一个空帧
						}
					} else {
						w.Trace("response: %s, len: %d", r, len(r))
						node.SendMessage(client, "", RESPONSE_OK, r) //回复REQ,因此要加上一个空帧
					}
				case COMMAND_SET, COMMAND_DEL, COMMAND_SCHEDULE: //key-value命令
					// 存到本地存储(同步)
					//回复结果(带信封, 否则找不到发送者), 因为是异步的, 可以先回复, 再做事
					if err := w.localStorage(cmd); err != nil {
						w.Debug("error: %s", err)
						node.SendMessage(client, "", RESPONSE_ERROR, err.Error()) //回复REQ,因此要加上一个空帧
					} else {
						node.SendMessage(client, "", RESPONSE_OK) //回复REQ,因此要加上一个空帧
					}

					// 发布(目标是跨IDC多点发布)
					publisher.SendMessage(cmd)

				case COMMAND_PUSH, COMMAND_TASK: //任务队列命令
					value := cmd[2:]
					if err := mqpool.Push(key, value); err == nil {
						w.Debug("push %s successful", key)
						node.SendMessage(client, "", RESPONSE_OK)
					} else {
						w.Debug("push %s failed: %s", key, err)
						node.SendMessage(client, "", RESPONSE_ERROR, err.Error())
					}
				case COMMAND_BTASK: //阻塞任务队列命令
					value := cmd[2:]
					taskId := ogoutils.NewShortUUID()
					value = append([]string{taskId}, value...) //放前面
					if err := mqpool.Push(key, value); err == nil {
						w.Debug("push block task %s successful, task id: %s", key, taskId)
						blockTasks[taskId] = make(chan string, 1)
						bto := time.Tick(BTASK_TIMEOUT)
						//go w.newBlocker(client)
						select {
						case <-bto: //超时
							w.Info("waiting time out")
							node.SendMessage(client, "", RESPONSE_ERROR)
						case result := <-blockTasks[taskId]:
							w.Debug("block task result: %s", result)
							node.SendMessage(client, "", RESPONSE_OK, result)
						}
						delete(blockTasks, taskId)
					} else {
						w.Debug("push %s failed: %s", key, err)
						node.SendMessage(client, "", RESPONSE_ERROR, err.Error())
					}
				case COMMAND_COMPLETE: // 完成阻塞任务
					if len(cmd) > 3 {
						if taskId := cmd[2]; taskId != "" {
							if _, ok := blockTasks[taskId]; ok {
								blockTasks[taskId] <- cmd[3]
								node.SendMessage(client, "", RESPONSE_OK)
							} else {
								node.SendMessage(client, "", RESPONSE_ERROR)
							}
						} else {
							node.SendMessage(client, "", RESPONSE_ERROR)
						}
					} else {
						node.SendMessage(client, "", RESPONSE_ERROR)
					}
				case COMMAND_POP:
					//cmd, _ := mqueuer.RecvMessage(0)
					if value, err := mqpool.Pop(key); err == nil {
						w.Debug("pop value from mqueue: %s", value)
						node.SendMessage(client, "", RESPONSE_OK, value) //回复REQ,因此要加上一个空帧
					} else if err.Error() == RESPONSE_NIL {
						w.Trace("pop %s nil: %s", key, err)
						node.SendMessage(client, "", RESPONSE_NIL) //没有内容,返回空
					} else {
						w.Trace("pop %s from mqueue failed: %s", key, err)
						node.SendMessage(client, "", RESPONSE_ERROR, err.Error()) //回复REQ,因此要加上一个空帧
					}
				default:
					// unknown action
					w.Info("unkown action: %s", act)
					node.SendMessage(client, "", RESPONSE_UNKNOWN)
				}

				// 回满血, 结束
				liveness = HEARTBEAT_LIVENESS
			} else if len(msg) == 1 {
				//  When we get a heartbeat message from the queue, it means the
				//  queue was (recently) alive, so reset our liveness indicator:
				if msg[0] == PPP_HEARTBEAT {
					liveness = HEARTBEAT_LIVENESS
				} else {
					w.Info("invalid message: %q", msg)
					node.SendMessage(client, "", RESPONSE_UNKNOWN)
				}
			} else {
				w.Info("invalid message: %q", msg)
				node.SendMessage(client, "", RESPONSE_UNKNOWN)
			}
			interval = INTERVAL_INIT
		} else {
			//  If the queue hasn't sent us heartbeats in a while, destroy the
			//  socket and reconnect. This is the simplest most brutal way of
			//  discarding any messages we might have sent in the meantime://
			liveness--
			if liveness == 0 {
				w.Debug("W: heartbeat failure, can't reach queue, reconnecting in %s", interval)

				time.Sleep(interval)

				if interval < INTERVAL_MAX { //每次重试都加大重试间隔
					interval = 2 * interval
				}
				// reconnect
				node, poller = w.connectQueue()
				liveness = HEARTBEAT_LIVENESS
			}
		}

		//  Send heartbeat to queue if it's time
		select {
		case <-heartbeat_at:
			if cycles > lastCycles {
				w.Trace("node%d worked cycles: %d", i, cycles)
				lastCycles = cycles
			}
			node.Send(PPP_HEARTBEAT, 0)
		default:
		}
	}
}

/* }}} */

/* {{{ func (w *OmqWorker) newBlocker(client string)
 * 异步阻塞节点
 */
func (w *OmqWorker) newBlocker(client string) {
	//node, _ := w.connectQueue()
	//w.Trace("block node for client: %q", client)

	//  Send out heartbeats at regular intervals
	heartbeat_at := time.Tick(5 * time.Second)

	//  Send heartbeat to queue if it's time
	select {
	case <-heartbeat_at:
		w.Debug("recv taskman response")
		//node.SendMessage(client, "", "BTASK")
	}
}

/* }}} */
<commit_msg>update<commit_after>package workers

import (
	"strings"
	"time"

	ogoutils "github.com/Odinman/ogo/utils"
	"github.com/Odinman/omq/utils"
	zmq "github.com/pebbe/zmq4"
)

/* {{{ func connectQueue() (*zmq.Socket, *zmq.Poller)
 *  Helper function that returns a new configured socket
 *  connected to the Paranoid Pirate queue
 */
func (w *OmqWorker) connectQueue() (*zmq.Socket, *zmq.Poller) {
	soc, _ := zmq.NewSocket(zmq.DEALER)
	soc.Connect("inproc://backend")

	//  Tell queue we're ready for work
	soc.Send(PPP_READY, 0)

	poller := zmq.NewPoller()
	poller.Add(soc, zmq.POLLIN)

	return soc, poller
}

/* }}} */

/* {{{ func (w *OmqWorker) newResponser(i int)
 * 回复节点
 */
func (w *OmqWorker) newResponser(i int) {
	node, poller := w.connectQueue()
	w.Trace("%d node ready", i)

	//  If liveness hits zero, queue is considered disconnected
	liveness := HEARTBEAT_LIVENESS
	interval := INTERVAL_INIT

	//  Send out heartbeats at regular intervals
	heartbeat_at := time.Tick(HEARTBEAT_INTERVAL)

	lastCycles := 0
	for cycles := 0; true; {
		sockets, err := poller.Poll(HEARTBEAT_INTERVAL)
		if err != nil {
			w.Error("polling error: %s", err)
			break //  Interrupted
		}

		if len(sockets) == 1 {
			//  Get message
			//  - n-part envelope + content -> request
			//  - 1-part HEARTBEAT -> heartbeat
			msg, err := node.RecvMessage(0)
			if err != nil {
				w.Error("recv error: %s", err)
				break //  Interrupted
			}

			client, cmd := utils.Unwrap(msg)
			if len(msg) >= 4 { //命令应该大于5帧(包含信封以及空帧)
				cycles++

				w.Trace("recv cmd: %s, from client: %q", cmd, client)

				act := strings.ToUpper(cmd[0])
				key := cmd[1]
				switch act {
				case COMMAND_GET, COMMAND_TIMING: //获取key内容
					if r, err := w.localGet(cmd); err != nil {
						w.Debug("error: %s", err)
						if err == ErrNil {
							node.SendMessage(client, "", RESPONSE_NIL) //没有内容,返回空
						} else {
							node.SendMessage(client, "", RESPONSE_ERROR, err.Error()) //回复REQ,因此要加上一个空帧
						}
					} else {
						w.Trace("response: %s, len: %d", r, len(r))
						node.SendMessage(client, "", RESPONSE_OK, r) //回复REQ,因此要加上一个空帧
					}
				case COMMAND_SET, COMMAND_DEL, COMMAND_SCHEDULE: //key-value命令
					// 存到本地存储(同步)
					//回复结果(带信封, 否则找不到发送者), 因为是异步的, 可以先回复, 再做事
					if err := w.localStorage(cmd); err != nil {
						w.Debug("error: %s", err)
						node.SendMessage(client, "", RESPONSE_ERROR, err.Error()) //回复REQ,因此要加上一个空帧
					} else {
						node.SendMessage(client, "", RESPONSE_OK) //回复REQ,因此要加上一个空帧
					}

					// 发布(目标是跨IDC多点发布)
					publisher.SendMessage(cmd)

				case COMMAND_PUSH, COMMAND_TASK: //任务队列命令
					value := cmd[2:]
					if err := mqpool.Push(key, value); err == nil {
						w.Debug("push %s successful", key)
						node.SendMessage(client, "", RESPONSE_OK)
					} else {
						w.Debug("push %s failed: %s", key, err)
						node.SendMessage(client, "", RESPONSE_ERROR, err.Error())
					}
				case COMMAND_BTASK: //阻塞任务队列命令
					value := cmd[2:]
					taskId := ogoutils.NewShortUUID()
					value = append([]string{taskId}, value...) //放前面
					if err := mqpool.Push(key, value); err == nil {
						w.Debug("push block task %s successful, task id: %s", key, taskId)
						blockTasks[taskId] = make(chan string, 1)
						bto := time.Tick(BTASK_TIMEOUT)
						//go w.newBlocker(client)
						select {
						case <-bto: //超时
							w.Info("waiting time out")
							node.SendMessage(client, "", RESPONSE_ERROR)
						case result := <-blockTasks[taskId]:
							w.Debug("block task result: %s", result)
							if result == "0" {
								node.SendMessage(client, "", RESPONSE_ERROR)
							} else {
								node.SendMessage(client, "", RESPONSE_OK, result)
							}
						}
						delete(blockTasks, taskId)
					} else {
						w.Debug("push %s failed: %s", key, err)
						node.SendMessage(client, "", RESPONSE_ERROR, err.Error())
					}
				case COMMAND_COMPLETE: // 完成阻塞任务
					if len(cmd) > 3 {
						if taskId := cmd[2]; taskId != "" {
							if _, ok := blockTasks[taskId]; ok {
								blockTasks[taskId] <- cmd[3]
								node.SendMessage(client, "", RESPONSE_OK)
							} else {
								node.SendMessage(client, "", RESPONSE_ERROR)
							}
						} else {
							node.SendMessage(client, "", RESPONSE_ERROR)
						}
					} else {
						node.SendMessage(client, "", RESPONSE_ERROR)
					}
				case COMMAND_POP:
					//cmd, _ := mqueuer.RecvMessage(0)
					if value, err := mqpool.Pop(key); err == nil {
						w.Debug("pop value from mqueue: %s", value)
						node.SendMessage(client, "", RESPONSE_OK, value) //回复REQ,因此要加上一个空帧
					} else if err.Error() == RESPONSE_NIL {
						w.Trace("pop %s nil: %s", key, err)
						node.SendMessage(client, "", RESPONSE_NIL) //没有内容,返回空
					} else {
						w.Trace("pop %s from mqueue failed: %s", key, err)
						node.SendMessage(client, "", RESPONSE_ERROR, err.Error()) //回复REQ,因此要加上一个空帧
					}
				default:
					// unknown action
					w.Info("unkown action: %s", act)
					node.SendMessage(client, "", RESPONSE_UNKNOWN)
				}

				// 回满血, 结束
				liveness = HEARTBEAT_LIVENESS
			} else if len(msg) == 1 {
				//  When we get a heartbeat message from the queue, it means the
				//  queue was (recently) alive, so reset our liveness indicator:
				if msg[0] == PPP_HEARTBEAT {
					liveness = HEARTBEAT_LIVENESS
				} else {
					w.Info("invalid message: %q", msg)
					node.SendMessage(client, "", RESPONSE_UNKNOWN)
				}
			} else {
				w.Info("invalid message: %q", msg)
				node.SendMessage(client, "", RESPONSE_UNKNOWN)
			}
			interval = INTERVAL_INIT
		} else {
			//  If the queue hasn't sent us heartbeats in a while, destroy the
			//  socket and reconnect. This is the simplest most brutal way of
			//  discarding any messages we might have sent in the meantime://
			liveness--
			if liveness == 0 {
				w.Debug("W: heartbeat failure, can't reach queue, reconnecting in %s", interval)

				time.Sleep(interval)

				if interval < INTERVAL_MAX { //每次重试都加大重试间隔
					interval = 2 * interval
				}
				// reconnect
				node, poller = w.connectQueue()
				liveness = HEARTBEAT_LIVENESS
			}
		}

		//  Send heartbeat to queue if it's time
		select {
		case <-heartbeat_at:
			if cycles > lastCycles {
				w.Trace("node%d worked cycles: %d", i, cycles)
				lastCycles = cycles
			}
			node.Send(PPP_HEARTBEAT, 0)
		default:
		}
	}
}

/* }}} */

/* {{{ func (w *OmqWorker) newBlocker(client string)
 * 异步阻塞节点
 */
func (w *OmqWorker) newBlocker(client string) {
	//node, _ := w.connectQueue()
	//w.Trace("block node for client: %q", client)

	//  Send out heartbeats at regular intervals
	heartbeat_at := time.Tick(5 * time.Second)

	//  Send heartbeat to queue if it's time
	select {
	case <-heartbeat_at:
		w.Debug("recv taskman response")
		//node.SendMessage(client, "", "BTASK")
	}
}

/* }}} */
<|endoftext|>


<commit_before>package eve

// Writer allows to write data to the EVE memory.
type Writer struct {
	d *Driver
}

func (w Writer) align32() {
	d := w.d
	m := (len(d.buf) + d.n) & 3
	if m == 0 {
		return
	}
	m = len(d.buf) + 4 - m
	if m > cap(d.buf) {
		m -= len(d.buf)
		d.flush()
	}
	d.buf = d.buf[:m]
}

func (w Writer) wr32(u uint32) {
	d := w.d
	if len(d.buf)+4 > cap(d.buf) {
		d.flush()
	}
	n := len(d.buf)
	d.buf = d.buf[:n+4]
	d.buf[n] = byte(u)
	d.buf[n+1] = byte(u >> 8)
	d.buf[n+2] = byte(u >> 16)
	d.buf[n+3] = byte(u >> 24)
}

func (w Writer) aw32(u uint32) {
	w.align32()
	w.wr32(u)
}

func (w Writer) Write32(s ...uint32) {
	w.align32()
	for _, u := range s {
		w.wr32(u)
	}
}

func (w Writer) Write(s []byte) (int, error) {
	d := w.d
	if len(s) >= cap(d.buf) {
		if len(d.buf) > 0 {
			d.flush()
		}
		d.dci.Write(s)
		d.n += len(s)
	} else {
		n := copy(d.buf[:cap(d.buf)], s)
		d.buf = d.buf[:len(d.buf)+n]
		if n < len(s) {
			d.flush()
			d.buf = d.buf[:len(s)-n]
			copy(d.buf, s[n:])
		}
	}
	// Write always succeeds. There is no case for infinite write transaction
	// so Driver.Err can be called after all writes.
	return len(s), nil
}

func (w Writer) WriteString(s string) (int, error) {
	d := w.d
	n := len(s)
	for len(s) != 0 {
		if len(d.buf) == cap(d.buf) {
			d.flush()
		}
		m := len(d.buf)
		c := copy(d.buf[m:cap(d.buf)], s)
		d.buf = d.buf[:m+c]
		s = s[c:]
	}
	return n, nil
}

// Close closes the write transaction and returns number of bytes written.
func (w Writer) Close() int {
	return w.d.end()
}
<commit_msg>display/eve: Fix bug in Writer.Write.<commit_after>package eve

// Writer allows to write data to the EVE memory.
type Writer struct {
	d *Driver
}

func (w Writer) align32() {
	d := w.d
	m := (len(d.buf) + d.n) & 3
	if m == 0 {
		return
	}
	m = len(d.buf) + 4 - m
	if m > cap(d.buf) {
		m -= len(d.buf)
		d.flush()
	}
	d.buf = d.buf[:m]
}

func (w Writer) wr8(b byte) {
	d := w.d
	if len(d.buf) == cap(d.buf) {
		d.flush()
	}
	n := len(d.buf)
	d.buf = d.buf[:n+1]
	d.buf[n] = b
}

func (w Writer) wr32(u uint32) {
	d := w.d
	if len(d.buf)+4 > cap(d.buf) {
		d.flush()
	}
	n := len(d.buf)
	d.buf = d.buf[:n+4]
	d.buf[n] = byte(u)
	d.buf[n+1] = byte(u >> 8)
	d.buf[n+2] = byte(u >> 16)
	d.buf[n+3] = byte(u >> 24)
}

func (w Writer) aw32(u uint32) {
	w.align32()
	w.wr32(u)
}

func (w Writer) Write8(s ...byte) {
	if len(s) == 1 {
		w.wr8(s[0])
		return
	}
	if len(s) == 0 {
		return
	}
	d := w.d
	if len(s) >= cap(d.buf) {
		if len(d.buf) > 0 {
			d.flush()
		}
		d.dci.Write(s)
		d.n += len(s)
		return
	}
	n := len(d.buf)
	m := copy(d.buf[n:cap(d.buf)], s)
	d.buf = d.buf[:n+m]
	if m < len(s) {
		s = s[m:]
		d.flush()
		d.buf = d.buf[:len(s)]
		copy(d.buf, s)
	}
}

func (w Writer) Write32(s ...uint32) {
	w.align32()
	for _, u := range s {
		w.wr32(u)
	}
}

func (w Writer) WriteInt(s ...int) {
	w.align32()
	for _, i := range s {
		w.wr32(uint32(i))
	}
}

func (w Writer) Write(s []byte) (int, error) {
	w.Write8(s...)
	return len(s), nil
	// BUG?: Write always succeeds. Rationale: there is no case for infinite
	// write transaction so Driver.Err can be called after all writes.
}

func (w Writer) wrs(s string) {
	d := w.d
	for len(s) != 0 {
		if len(d.buf) == cap(d.buf) {
			d.flush()
		}
		n := len(d.buf)
		c := copy(d.buf[n:cap(d.buf)], s)
		d.buf = d.buf[:n+c]
		s = s[c:]
	}
}

func (w Writer) WriteString(s string) (int, error) {
	w.wrs(s)
	return len(s), nil
	// BUG?: WriteString always succeeds. Rationale: there is no case for
	// infinite write transaction so Driver.Err can be called after all writes.
}

// Close closes the write transaction and returns number of bytes written.
func (w Writer) Close() int {
	return w.d.end()
}
<|endoftext|>


<commit_before>package gpio

type Mode byte

const (
	In     Mode = 0
	InOut  Mode = 1
	Discon Mode = 2
	Out    Mode = 3
)

type Pull byte

const (
	NoPull   Pull = 0
	PullDown Pull = 1
	PullUp   Pull = 3
)

type Config struct {
	Mode Mode
	Pull Pull
}

func (p *Port) SetupPin(n int, cfg *Config) {
	p.pincnf[n].Store(uint32(cfg.Pull)<<2 | uint32(cfg.Mode))
}

// Setup configures pins.
func (p *Port) Setup(pins Pins, cfg *Config) {
	for n := 0; n < 32; n++ {
		if pins&(1<<uint(n)) != 0 {
			p.SetupPin(n, cfg)
		}
	}
}
<commit_msg>nrf5/hal/gpio: More configuration options.<commit_after>package gpio

type Mode byte

const (
	In     Mode = 0
	InOut  Mode = 1
	Discon Mode = 2
	Out    Mode = 3
)

type Pull byte

const (
	NoPull   Pull = 0
	PullDown Pull = 1
	PullUp   Pull = 3
)

type Drive byte

const (
	S0S1 Drive = 0 // Standard 0, standard 1.
	H0S1 Drive = 1 // High drive 0, standard 1.
	S0H1 Drive = 2 // Standard 0, high drive 1.
	H0H1 Drive = 3 // High drive 0, high drive 1.
	D0S1 Drive = 4 // Disconnect 0, standard 1.
	D0H1 Drive = 5 // Disconnect 0, high drive 1.
	S0D1 Drive = 6 // Standard 0, disconnect 1.
	H0D1 Drive = 7 // High drive 0, disconnect 1.
)

type Sense byte

const (
	NoSense   Sense = 0
	SenseHigh Sense = 2
	SenseLow  Sense = 3
)

type Config struct {
	Mode  Mode
	Pull  Pull
	Drive Drive
	Sense Sense
}

func (p *Port) SetupPin(n int, cfg *Config) {
	p.pincnf[n].Store(
		uint32(cfg.Sense)<<16 | uint32(cfg.Drive)<<8 | uint32(cfg.Pull)<<2 |
			uint32(cfg.Mode),
	)
}

// Setup configures pins.
func (p *Port) Setup(pins Pins, cfg *Config) {
	for n := 0; n < 32; n++ {
		if pins&(1<<uint(n)) != 0 {
			p.SetupPin(n, cfg)
		}
	}
}

// SetDirIn allows a fast change of direction to input for specified pins.
func (p *Port) SetDirIn(pins Pins) {
	p.dirset.Store(uint32(pins))
}

// SetDirOut allows afast change of direction to output for specified pins.
func (p *Port) SetDirOut(pins Pins) {
	p.dirclr.Store(uint32(pins))
}
<|endoftext|>


<commit_before>// Package ppipwm allows to produce PWM signal using PPI, TIMER and GPIOTE
// peripherals.
package ppipwm

import (
	"nrf5/hal/gpio"
	"nrf5/hal/gpiote"
	"nrf5/hal/ppi"
	"nrf5/hal/timer"
)

// Toggle implementats three channel PWM usind GPIOTE OUT task configured in
// toggle mode. It is designed specifically for nRF51 chips because of
// limitations of their GPIOTE peripheral. Better solutions exists for nRF52.
//
// To produce desired PWM waveform Toggle configures a timer compare register,
// using simple algorithm, that usually generates glitches when used on working
// PWM channel. After that, PWM works without any CPU intervention and produces
// proper waveform until next duty cycle change. This is fine to drive LEDs but
// can cause trubles in case of receivers that rely on stable PWM frequency or
// phase (eg. some servos). Correct implementation is difficult and quite
// expensive in case of nRF51.
//
// Toggle cannot be used concurently by multiple gorutines without proper
// synchronisation.
type Toggle struct {
	t  *timer.Periph
	gc [3]gpiote.Chan
}

// MakeToggle returns configured Toggle implementation of PPI based PWM using
// timer t.
func MakeToggle(t *timer.Periph) Toggle {
	t.Task(timer.STOP).Trigger()
	t.StoreSHORTS(timer.COMPARE3_CLEAR)
	return Toggle{t: t}
}

// NewToggle provides convenient way to create heap allocated Toggle struct. See
// MakeToggle for more information.
func NewToggle(t *timer.Periph) *Toggle {
	pwm := new(Toggle)
	*pwm = MakeToggle(t)
	return pwm
}

// SetFreq sets prescaler to 2^log2pre and period to periodus microseconds. It
// allows to configure a duty cycle with resolution = 16 * periodus / 2^log2pre.
// SetFreq returns (resolution-1), which is a value that should be passed to
// SetDuty/SetInvDuty to produce PWM with 100% duty cycle.
func (pwm *Toggle) SetFreq(log2pre, periodus int) int {
	if uint(log2pre) > 9 {
		panic("pwm: bad prescaler")
	}
	if periodus < 1 {
		panic("pwm: bad period")
	}
	t := pwm.t
	t.StorePRESCALER(log2pre)
	div := uint32(1) << uint(log2pre)
	max := 16*uint32(periodus)/div - 1
	if max > 0xFFFF {
		panic("pwm: max>0xFFFF")
	}
	t.StoreCC(3, max)
	return int(max)
}

// Max returns a value that corresponds to 100% PWM duty cycle. See SetFreq for
// more information.
func (pwm *Toggle) Max() int {
	return int(pwm.t.LoadCC(3))
}

func checkOutput(n int) {
	if uint(n) > 2 {
		panic("pwm: bad output")
	}
}

// Setup setups n-th of three PWM channels. Each PWM channel uses one GPIOTE
// channel and two PPI channels.
func (pwm *Toggle) Setup(n int, pin gpio.Pin, gc gpiote.Chan, pc0, pc1 ppi.Chan) {
	checkOutput(n)
	pin.Clear()
	pin.Setup(gpio.ModeOut)
	gc.Setup(pin, 0)
	pwm.gc[n] = gc
	t := pwm.t
	pc0.SetEEP(t.Event(timer.COMPARE(n)))
	pc0.SetTEP(gc.OUT().Task())
	pc0.Enable()
	pc1.SetEEP(t.Event(timer.COMPARE(3)))
	pc1.SetTEP(gc.OUT().Task())
	pc1.Enable()
}

// SetVal sets a duty cycle for n-th PWM channel.
func (pwm *Toggle) SetVal(n, val int) {
	checkOutput(n)
	gc := pwm.gc[n]
	t := pwm.t
	pin, _ := gc.Config()
	t.Task(timer.STOP).Trigger()
	switch {
	case val <= 0:
		pin.Clear()
		gc.Setup(pin, 0)
		return
	case val >= pwm.Max():
		pin.Set()
		gc.Setup(pin, 0)
		return
	}
	cfg := gpiote.ModeTask | gpiote.PolarityToggle
	t.Task(timer.CAPTURE(n)).Trigger()
	cnt := int(t.LoadCC(n))
	if cnt < val {
		gc.Setup(pin, cfg|gpiote.OutInitHigh)
	} else {
		gc.Setup(pin, cfg|gpiote.OutInitLow)
	}
	t.StoreCC(n, uint32(val))
	t.Task(timer.START).Trigger()
}

// SetInvDuty sets a duty cycle for n-th PWM channel. It produces inverted
// waveform.
func (pwm *Toggle) SetInvVal(n, val int) {
	checkOutput(n)
	gc := pwm.gc[n]
	t := pwm.t
	pin, _ := gc.Config()
	t.Task(timer.STOP).Trigger()
	switch {
	case val <= 0:
		pin.Set()
		gc.Setup(pin, 0)
		return
	case val >= pwm.Max():
		pin.Clear()
		gc.Setup(pin, 0)
		return
	}
	cfg := gpiote.ModeTask | gpiote.PolarityToggle
	t.Task(timer.CAPTURE(n)).Trigger()
	cnt := int(t.LoadCC(n))
	if cnt < val {
		gc.Setup(pin, cfg|gpiote.OutInitLow)
	} else {
		gc.Setup(pin, cfg|gpiote.OutInitHigh)
	}
	t.StoreCC(n, uint32(val))
	t.Task(timer.START).Trigger()
}
<commit_msg>nrf5/ppipwm: Better description of SetFreq method.<commit_after>// Package ppipwm allows to produce PWM signal using PPI, TIMER and GPIOTE
// peripherals.
package ppipwm

import (
	"nrf5/hal/gpio"
	"nrf5/hal/gpiote"
	"nrf5/hal/ppi"
	"nrf5/hal/timer"
)

// Toggle implementats three channel PWM usind GPIOTE OUT task configured in
// toggle mode. It is designed specifically for nRF51 chips because of
// limitations of their GPIOTE peripheral. Better solutions exists for nRF52.
//
// To produce desired PWM waveform Toggle configures a timer compare register,
// using simple algorithm, that usually generates glitches when used on working
// PWM channel. After that, PWM works without any CPU intervention and produces
// proper waveform until next duty cycle change. This is fine to drive LEDs but
// can cause trubles in case of receivers that rely on stable PWM frequency or
// phase (eg. some servos). Correct implementation is difficult and quite
// expensive in case of nRF51.
//
// Toggle cannot be used concurently by multiple gorutines without proper
// synchronisation.
type Toggle struct {
	t  *timer.Periph
	gc [3]gpiote.Chan
}

// MakeToggle returns configured Toggle implementation of PPI based PWM using
// timer t.
func MakeToggle(t *timer.Periph) Toggle {
	t.Task(timer.STOP).Trigger()
	t.StoreSHORTS(timer.COMPARE3_CLEAR)
	return Toggle{t: t}
}

// NewToggle provides convenient way to create heap allocated Toggle struct. See
// MakeToggle for more information.
func NewToggle(t *timer.Periph) *Toggle {
	pwm := new(Toggle)
	*pwm = MakeToggle(t)
	return pwm
}

// SetFreq sets timer prescaler to 2^log2pre and period to periodus microseconds
// (log2pre must be in range [0..9]). It allows to configure a duty cycle with
// a resolution = 16 * periodus / 2^log2pre. Toggle uses timer in 16-bit mode so
// the resolution must be <= 65536. SetFreq returns (resolution-1), which is a
// value that should be passed to SetDuty/SetInvDuty to produce PWM with 100%
// duty cycle.
func (pwm *Toggle) SetFreq(log2pre, periodus int) int {
	if uint(log2pre) > 9 {
		panic("pwm: bad prescaler")
	}
	if periodus < 1 {
		panic("pwm: bad period")
	}
	t := pwm.t
	t.StorePRESCALER(log2pre)
	div := uint32(1) << uint(log2pre)
	max := 16*uint32(periodus)/div - 1 // 16 MHz * period / div - 1
	if max > 0xFFFF {
		panic("pwm: max>0xFFFF")
	}
	t.StoreCC(3, max)
	return int(max)
}

// Max returns a value that corresponds to 100% PWM duty cycle. See SetFreq for
// more information.
func (pwm *Toggle) Max() int {
	return int(pwm.t.LoadCC(3))
}

func checkOutput(n int) {
	if uint(n) > 2 {
		panic("pwm: bad output")
	}
}

// Setup setups n-th of three PWM channels. Each PWM channel uses one GPIOTE
// channel and two PPI channels.
func (pwm *Toggle) Setup(n int, pin gpio.Pin, gc gpiote.Chan, pc0, pc1 ppi.Chan) {
	checkOutput(n)
	pin.Clear()
	pin.Setup(gpio.ModeOut)
	gc.Setup(pin, 0)
	pwm.gc[n] = gc
	t := pwm.t
	pc0.SetEEP(t.Event(timer.COMPARE(n)))
	pc0.SetTEP(gc.OUT().Task())
	pc0.Enable()
	pc1.SetEEP(t.Event(timer.COMPARE(3)))
	pc1.SetTEP(gc.OUT().Task())
	pc1.Enable()
}

// SetVal sets a duty cycle for n-th PWM channel.
func (pwm *Toggle) SetVal(n, val int) {
	checkOutput(n)
	gc := pwm.gc[n]
	t := pwm.t
	pin, _ := gc.Config()
	t.Task(timer.STOP).Trigger()
	switch {
	case val <= 0:
		pin.Clear()
		gc.Setup(pin, 0)
		return
	case val >= pwm.Max():
		pin.Set()
		gc.Setup(pin, 0)
		return
	}
	cfg := gpiote.ModeTask | gpiote.PolarityToggle
	t.Task(timer.CAPTURE(n)).Trigger()
	cnt := int(t.LoadCC(n))
	if cnt < val {
		gc.Setup(pin, cfg|gpiote.OutInitHigh)
	} else {
		gc.Setup(pin, cfg|gpiote.OutInitLow)
	}
	t.StoreCC(n, uint32(val))
	t.Task(timer.START).Trigger()
}

// SetInvDuty sets a duty cycle for n-th PWM channel. It produces inverted
// waveform.
func (pwm *Toggle) SetInvVal(n, val int) {
	checkOutput(n)
	gc := pwm.gc[n]
	t := pwm.t
	pin, _ := gc.Config()
	t.Task(timer.STOP).Trigger()
	switch {
	case val <= 0:
		pin.Set()
		gc.Setup(pin, 0)
		return
	case val >= pwm.Max():
		pin.Clear()
		gc.Setup(pin, 0)
		return
	}
	cfg := gpiote.ModeTask | gpiote.PolarityToggle
	t.Task(timer.CAPTURE(n)).Trigger()
	cnt := int(t.LoadCC(n))
	if cnt < val {
		gc.Setup(pin, cfg|gpiote.OutInitLow)
	} else {
		gc.Setup(pin, cfg|gpiote.OutInitHigh)
	}
	t.StoreCC(n, uint32(val))
	t.Task(timer.START).Trigger()
}
<|endoftext|>


<commit_before>// Package besticon includes functions
// finding icons for a given web site.
package besticon

import (
	"bytes"
	"crypto/sha1"
	"errors"
	"fmt"
	"image"
	"io/ioutil"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"sort"
	"strings"

	// Load support for common image formats.
	_ "code.google.com/p/go.image/bmp"
	_ "code.google.com/p/go.image/tiff"
	_ "code.google.com/p/go.image/webp"
	_ "github.com/mat/besticon/ico"
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"

	"code.google.com/p/go.net/html/charset"
	"code.google.com/p/go.net/publicsuffix"
	"github.com/PuerkitoBio/goquery"
)

// Icon holds icon information.
type Icon struct {
	URL     string `json:"url"`
	Width   int    `json:"width"`
	Height  int    `json:"height"`
	Format  string `json:"format"`
	Bytes   int    `json:"bytes"`
	Error   error  `json:"error"`
	Sha1sum string `json:"sha1sum"`
}

// FetchBestIcon takes a siteURL and returns the icon with
// the largest dimensions for this site or an error.
func FetchBestIcon(siteURL string) (*Icon, error) {
	return fetchBestIconWithClient(siteURL, &http.Client{})
}

func fetchBestIconWithClient(siteURL string, c *http.Client) (*Icon, error) {
	icons, e := fetchIconsWithClient(siteURL, c)
	if e != nil {
		return nil, e
	}

	if len(icons) < 1 {
		return nil, errors.New("besticon: no icons found for site")
	}

	best := icons[0]
	return &best, nil
}

// FetchIcons takes a siteURL and returns all icons for this site
// or an error.
func FetchIcons(siteURL string) ([]Icon, error) {
	return fetchIconsWithClient(siteURL, &http.Client{})
}

// fetchIconsWithClient modifies c's checkRedirect and Jar!
func fetchIconsWithClient(siteURL string, c *http.Client) ([]Icon, error) {
	configureClient(c)

	siteURL = strings.TrimSpace(siteURL)
	if !strings.HasPrefix(siteURL, "http") {
		siteURL = "http://" + siteURL
	}

	html, url, e := fetchHTML(siteURL, c)
	if e != nil {
		return nil, e
	}

	links, e := assembleIconLinks(url, html)
	if e != nil {
		return nil, e
	}

	icons := fetchAllIcons(links, c)
	icons = rejectBrokenIcons(icons)

	// Order when finished: (width/height, bytes, url)
	sort.Stable(byURL(icons))
	sort.Stable(byBytes(icons))
	sort.Stable(sort.Reverse(byWidthHeight(icons)))

	return icons, nil
}

func fetchHTML(url string, c *http.Client) ([]byte, *url.URL, error) {
	r, e := get(c, url)
	if e != nil {
		return nil, nil, e
	}

	if !(r.StatusCode >= 200 && r.StatusCode < 300) {
		return nil, nil, errors.New("besticon: not found")
	}

	b, e := ioutil.ReadAll(r.Body)
	if e != nil {
		return nil, nil, e
	}
	defer r.Body.Close()
	if len(b) == 0 {
		return nil, nil, errors.New("besticon: empty response")
	}

	reader := bytes.NewReader(b)
	contentType := r.Header.Get("Content-Type")
	utf8reader, e := charset.NewReader(reader, contentType)
	if e != nil {
		return nil, nil, e
	}
	utf8bytes, e := ioutil.ReadAll(utf8reader)
	if e != nil {
		return nil, nil, e
	}

	return utf8bytes, r.Request.URL, nil
}

var iconPaths = []string{
	"/favicon.ico",
	"/apple-touch-icon.png",
	"/apple-touch-icon-precomposed.png",
}

type empty struct{}

func assembleIconLinks(siteURL *url.URL, html []byte) ([]string, error) {
	links := make(map[string]empty)

	// Add common, hard coded icon paths
	for _, path := range iconPaths {
		links[urlFromBase(siteURL, path)] = empty{}
	}

	// Add icons found in page
	urls, e := findIcons(html)
	if e != nil {
		return nil, e
	}
	for _, url := range urls {
		url, e := absoluteURL(siteURL, url)
		if e == nil {
			links[url] = empty{}
		}
	}

	// Turn unique keys into array
	result := []string{}
	for u := range links {
		result = append(result, u)
	}
	sort.Strings(result)

	return result, nil
}

var csspaths = strings.Join([]string{
	"link[rel='icon']",
	"link[rel='shortcut icon']",
	"link[rel='apple-touch-icon']",
	"link[rel='apple-touch-icon-precomposed']",
}, ", ")

var errParseHTML = errors.New("besticon: could not parse html")

func findIcons(html []byte) ([]string, error) {
	doc, e := goquery.NewDocumentFromReader(bytes.NewReader(html))
	if e != nil || doc == nil {
		return nil, errParseHTML
	}

	hits := []string{}
	doc.Find(csspaths).Each(func(i int, s *goquery.Selection) {
		href, ok := s.Attr("href")
		if ok && href != "" {
			hits = append(hits, href)
		}
	})

	return hits, nil
}

func fetchAllIcons(urls []string, c *http.Client) []Icon {
	ch := make(chan Icon)

	for _, u := range urls {
		go func(u string) { ch <- fetchIconDetails(u, c) }(u)
	}

	icons := []Icon{}
	for range urls {
		icon := <-ch
		icons = append(icons, icon)
	}
	return icons
}

func fetchIconDetails(url string, c *http.Client) Icon {
	i := Icon{URL: url}

	response, e := get(c, url)
	if e != nil {
		i.Error = e
		return i
	}

	b, e := ioutil.ReadAll(response.Body)
	if e != nil {
		i.Error = e
		return i
	}
	defer response.Body.Close()

	cfg, format, e := image.DecodeConfig(bytes.NewReader(b))
	if e != nil {
		i.Error = fmt.Errorf("besticon: unknown image format: %s", e)
		return i
	}

	i.Width = cfg.Width
	i.Height = cfg.Height
	i.Format = format
	i.Bytes = len(b)
	i.Sha1sum = sha1Sum(b)

	return i
}

func get(client *http.Client, url string) (*http.Response, error) {
	req, e := http.NewRequest("GET", url, nil)
	if e != nil {
		return nil, e
	}

	setDefaultHeaders(req)
	return client.Do(req)
}

func setDefaultHeaders(req *http.Request) {
	req.Header.Set("Accept", "*/*")
	req.Header.Set("User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_7_3) AppleWebKit/534.55.3 (KHTML, like Gecko) Version/5.1.3 Safari/534.53.10")
}

func configureClient(c *http.Client) {
	c.Jar = mustInitCookieJar()
	c.CheckRedirect = checkRedirect
}

func mustInitCookieJar() *cookiejar.Jar {
	options := cookiejar.Options{
		PublicSuffixList: publicsuffix.List,
	}
	jar, e := cookiejar.New(&options)
	if e != nil {
		panic(e)
	}

	return jar
}

func checkRedirect(req *http.Request, via []*http.Request) error {
	setDefaultHeaders(req)

	if len(via) >= 10 {
		return errors.New("stopped after 10 redirects")
	}
	return nil
}

func absoluteURL(baseURL *url.URL, path string) (string, error) {
	url, e := url.Parse(path)
	if e != nil {
		return "", e
	}

	url.Scheme = baseURL.Scheme
	if url.Host == "" {
		url.Host = baseURL.Host
	}
	return url.String(), nil
}

func urlFromBase(baseURL *url.URL, path string) string {
	url := *baseURL
	url.Path = path
	return url.String()
}

func rejectBrokenIcons(icons []Icon) []Icon {
	result := []Icon{}
	for _, img := range icons {
		if img.Error == nil && (img.Width > 1 && img.Height > 1) {
			result = append(result, img)
		}
	}
	return result
}

func sha1Sum(b []byte) string {
	hash := sha1.New()
	hash.Write(b)
	bs := hash.Sum(nil)
	return fmt.Sprintf("%x", bs)
}

func (i *Icon) ImgWidth() int {
	return i.Width / 2.0
}

type byWidthHeight []Icon

func (a byWidthHeight) Len() int      { return len(a) }
func (a byWidthHeight) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a byWidthHeight) Less(i, j int) bool {
	return (a[i].Width < a[j].Width) || (a[i].Height < a[j].Height)
}

type byBytes []Icon

func (a byBytes) Len() int           { return len(a) }
func (a byBytes) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a byBytes) Less(i, j int) bool { return (a[i].Bytes < a[j].Bytes) }

type byURL []Icon

func (a byURL) Len() int           { return len(a) }
func (a byURL) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a byURL) Less(i, j int) bool { return (a[i].URL < a[j].URL) }
<commit_msg>Set 60s timeout on outgoing HTTP requests.<commit_after>// Package besticon includes functions
// finding icons for a given web site.
package besticon

import (
	"bytes"
	"crypto/sha1"
	"errors"
	"fmt"
	"image"
	"io/ioutil"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"sort"
	"strings"
	"time"

	// Load support for common image formats.
	_ "code.google.com/p/go.image/bmp"
	_ "code.google.com/p/go.image/tiff"
	_ "code.google.com/p/go.image/webp"
	_ "github.com/mat/besticon/ico"
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"

	"code.google.com/p/go.net/html/charset"
	"code.google.com/p/go.net/publicsuffix"
	"github.com/PuerkitoBio/goquery"
)

// Icon holds icon information.
type Icon struct {
	URL     string `json:"url"`
	Width   int    `json:"width"`
	Height  int    `json:"height"`
	Format  string `json:"format"`
	Bytes   int    `json:"bytes"`
	Error   error  `json:"error"`
	Sha1sum string `json:"sha1sum"`
}

// FetchBestIcon takes a siteURL and returns the icon with
// the largest dimensions for this site or an error.
func FetchBestIcon(siteURL string) (*Icon, error) {
	return fetchBestIconWithClient(siteURL, &http.Client{})
}

func fetchBestIconWithClient(siteURL string, c *http.Client) (*Icon, error) {
	icons, e := fetchIconsWithClient(siteURL, c)
	if e != nil {
		return nil, e
	}

	if len(icons) < 1 {
		return nil, errors.New("besticon: no icons found for site")
	}

	best := icons[0]
	return &best, nil
}

// FetchIcons takes a siteURL and returns all icons for this site
// or an error.
func FetchIcons(siteURL string) ([]Icon, error) {
	c := &http.Client{Timeout: 60 * time.Second}
	return fetchIconsWithClient(siteURL, c)
}

// fetchIconsWithClient modifies c's checkRedirect and Jar!
func fetchIconsWithClient(siteURL string, c *http.Client) ([]Icon, error) {
	configureClient(c)

	siteURL = strings.TrimSpace(siteURL)
	if !strings.HasPrefix(siteURL, "http") {
		siteURL = "http://" + siteURL
	}

	html, url, e := fetchHTML(siteURL, c)
	if e != nil {
		return nil, e
	}

	links, e := assembleIconLinks(url, html)
	if e != nil {
		return nil, e
	}

	icons := fetchAllIcons(links, c)
	icons = rejectBrokenIcons(icons)

	// Order when finished: (width/height, bytes, url)
	sort.Stable(byURL(icons))
	sort.Stable(byBytes(icons))
	sort.Stable(sort.Reverse(byWidthHeight(icons)))

	return icons, nil
}

func fetchHTML(url string, c *http.Client) ([]byte, *url.URL, error) {
	r, e := get(c, url)
	if e != nil {
		return nil, nil, e
	}

	if !(r.StatusCode >= 200 && r.StatusCode < 300) {
		return nil, nil, errors.New("besticon: not found")
	}

	b, e := ioutil.ReadAll(r.Body)
	if e != nil {
		return nil, nil, e
	}
	defer r.Body.Close()
	if len(b) == 0 {
		return nil, nil, errors.New("besticon: empty response")
	}

	reader := bytes.NewReader(b)
	contentType := r.Header.Get("Content-Type")
	utf8reader, e := charset.NewReader(reader, contentType)
	if e != nil {
		return nil, nil, e
	}
	utf8bytes, e := ioutil.ReadAll(utf8reader)
	if e != nil {
		return nil, nil, e
	}

	return utf8bytes, r.Request.URL, nil
}

var iconPaths = []string{
	"/favicon.ico",
	"/apple-touch-icon.png",
	"/apple-touch-icon-precomposed.png",
}

type empty struct{}

func assembleIconLinks(siteURL *url.URL, html []byte) ([]string, error) {
	links := make(map[string]empty)

	// Add common, hard coded icon paths
	for _, path := range iconPaths {
		links[urlFromBase(siteURL, path)] = empty{}
	}

	// Add icons found in page
	urls, e := findIcons(html)
	if e != nil {
		return nil, e
	}
	for _, url := range urls {
		url, e := absoluteURL(siteURL, url)
		if e == nil {
			links[url] = empty{}
		}
	}

	// Turn unique keys into array
	result := []string{}
	for u := range links {
		result = append(result, u)
	}
	sort.Strings(result)

	return result, nil
}

var csspaths = strings.Join([]string{
	"link[rel='icon']",
	"link[rel='shortcut icon']",
	"link[rel='apple-touch-icon']",
	"link[rel='apple-touch-icon-precomposed']",
}, ", ")

var errParseHTML = errors.New("besticon: could not parse html")

func findIcons(html []byte) ([]string, error) {
	doc, e := goquery.NewDocumentFromReader(bytes.NewReader(html))
	if e != nil || doc == nil {
		return nil, errParseHTML
	}

	hits := []string{}
	doc.Find(csspaths).Each(func(i int, s *goquery.Selection) {
		href, ok := s.Attr("href")
		if ok && href != "" {
			hits = append(hits, href)
		}
	})

	return hits, nil
}

func fetchAllIcons(urls []string, c *http.Client) []Icon {
	ch := make(chan Icon)

	for _, u := range urls {
		go func(u string) { ch <- fetchIconDetails(u, c) }(u)
	}

	icons := []Icon{}
	for range urls {
		icon := <-ch
		icons = append(icons, icon)
	}
	return icons
}

func fetchIconDetails(url string, c *http.Client) Icon {
	i := Icon{URL: url}

	response, e := get(c, url)
	if e != nil {
		i.Error = e
		return i
	}

	b, e := ioutil.ReadAll(response.Body)
	if e != nil {
		i.Error = e
		return i
	}
	defer response.Body.Close()

	cfg, format, e := image.DecodeConfig(bytes.NewReader(b))
	if e != nil {
		i.Error = fmt.Errorf("besticon: unknown image format: %s", e)
		return i
	}

	i.Width = cfg.Width
	i.Height = cfg.Height
	i.Format = format
	i.Bytes = len(b)
	i.Sha1sum = sha1Sum(b)

	return i
}

func get(client *http.Client, url string) (*http.Response, error) {
	req, e := http.NewRequest("GET", url, nil)
	if e != nil {
		return nil, e
	}

	setDefaultHeaders(req)
	return client.Do(req)
}

func setDefaultHeaders(req *http.Request) {
	req.Header.Set("Accept", "*/*")
	req.Header.Set("User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_7_3) AppleWebKit/534.55.3 (KHTML, like Gecko) Version/5.1.3 Safari/534.53.10")
}

func configureClient(c *http.Client) {
	c.Jar = mustInitCookieJar()
	c.CheckRedirect = checkRedirect
}

func mustInitCookieJar() *cookiejar.Jar {
	options := cookiejar.Options{
		PublicSuffixList: publicsuffix.List,
	}
	jar, e := cookiejar.New(&options)
	if e != nil {
		panic(e)
	}

	return jar
}

func checkRedirect(req *http.Request, via []*http.Request) error {
	setDefaultHeaders(req)

	if len(via) >= 10 {
		return errors.New("stopped after 10 redirects")
	}
	return nil
}

func absoluteURL(baseURL *url.URL, path string) (string, error) {
	url, e := url.Parse(path)
	if e != nil {
		return "", e
	}

	url.Scheme = baseURL.Scheme
	if url.Host == "" {
		url.Host = baseURL.Host
	}
	return url.String(), nil
}

func urlFromBase(baseURL *url.URL, path string) string {
	url := *baseURL
	url.Path = path
	return url.String()
}

func rejectBrokenIcons(icons []Icon) []Icon {
	result := []Icon{}
	for _, img := range icons {
		if img.Error == nil && (img.Width > 1 && img.Height > 1) {
			result = append(result, img)
		}
	}
	return result
}

func sha1Sum(b []byte) string {
	hash := sha1.New()
	hash.Write(b)
	bs := hash.Sum(nil)
	return fmt.Sprintf("%x", bs)
}

func (i *Icon) ImgWidth() int {
	return i.Width / 2.0
}

type byWidthHeight []Icon

func (a byWidthHeight) Len() int      { return len(a) }
func (a byWidthHeight) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a byWidthHeight) Less(i, j int) bool {
	return (a[i].Width < a[j].Width) || (a[i].Height < a[j].Height)
}

type byBytes []Icon

func (a byBytes) Len() int           { return len(a) }
func (a byBytes) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a byBytes) Less(i, j int) bool { return (a[i].Bytes < a[j].Bytes) }

type byURL []Icon

func (a byURL) Len() int           { return len(a) }
func (a byURL) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a byURL) Less(i, j int) bool { return (a[i].URL < a[j].URL) }
<|endoftext|>


<commit_before>package main

import (
	"os"
	"testing"
	"time"
	"bytes"
	"io"
)

// change traq env to use our fixtures
func WithFakeEnv(block func()) {
  oldEnv := os.Getenv("TRAQ_DATA_DIR")
  path, _ := os.Getwd()
  os.Setenv("TRAQ_DATA_DIR", path + "/fixtures")

  block()

  os.Setenv("TRAQ_DATA_DIR", oldEnv)
}

// capture output written to os.Stdout and return it
func CaptureStdout(block func()) string {
  old := os.Stdout // keep backup of the real stdout
  r, w, _ := os.Pipe()
  os.Stdout = w

  block()

  outC := make(chan string)
  // copy the output in a separate goroutine so printing can't block indefinitely
  go func() {
      var buf bytes.Buffer
      io.Copy(&buf, r)
      outC <- buf.String()
  }()

  // back to normal state
  w.Close()
  os.Stdout = old

  return <-outC
}

func TestDatesInMonth(t *testing.T) {
  dates := DatesInMonth(1986, 9)

  if len(dates) != 30 {
    t.Errorf("expected 30 days in Sep 1986, got %v", len(dates))
  }

  if dates[0].Weekday() != time.Monday {
    t.Errorf("Started on a Monday, got %v", dates[0].Weekday())
  }

  if dates[len(dates) - 1].Weekday() != time.Tuesday {
    t.Errorf("Ended on a Tuesday, got %v", dates[len(dates) - 1].Weekday())
  }
}

func TestPrintDate(t *testing.T) {
  out := CaptureStdout(func() {
    WithFakeEnv(func() {
      PrintDate("example", time.Date(1986, 9, 3, 0, 0, 0, 0, time.UTC))
    })
  })

  expected :=
`Wed Sep 03 20:00:00 +0100 1986;#birth;comment
Wed Sep 03 21:45:33 +0100 1986;#chillout;
Wed Sep 03 23:24:49 +0100 1986;stop;
%%
`
  if out != expected {
    t.Errorf("unexpected PrintDate output. Expected '%v' got '%v'", expected, out)
  }
}

func TestFilePath(t *testing.T) {
	var path string = FilePath("example", time.Date(1986, 9, 3, 0, 0, 0, 0, time.UTC))

	if path != os.Getenv("TRAQ_DATA_DIR")+"/example/1986/1986-09-03" {
		t.Errorf("FilePath = %v, want %v", path, os.Getenv("TRAQ_DATA_DIR")+"/example/1986/1986-09-03")
	}
}

func TestEmptySumFile(t *testing.T) {
	content := []string{""}
	var summed, error = SumFile(content)

	if error == nil {
		var total, ok = summed["#work"]
		if ok {
			t.Errorf("summed['#work'] = %v, should not exist", total)
		}
	} else {
		t.Errorf("parsing error %v", error)
	}
}

func TestSimpleSumFile(t *testing.T) {
	content := []string{
		"Mon Oct 28 21:45:33 +0100 2013;#work;",
		"Mon Oct 28 23:24:49 +0100 2013;stop;",
	}
	var summed, error = SumFile(content)

	if error == nil {
		var total, ok = summed["#work"]
		if total != 5956 || !ok {
			t.Errorf("summed['#work'] = %v, want %v", total, 5956)
		}
	} else {
		t.Errorf("parsing error %v", error)
	}
}

func TestNoStopSumFile(t *testing.T) {
	content := []string{
		"Mon Oct 28 20:00:00 +0100 2013;#play;",
		"Mon Oct 28 21:45:33 +0100 2013;#work;",
		"Mon Oct 28 23:24:49 +0100 2013;stop;",
	}
	var summed, error = SumFile(content)

	if error == nil {
		var total, ok = summed["#play"]
		if total != 6333 || !ok {
			t.Errorf("summed['#play'] = %v, want %v", total, 6333)
		}
		total, ok = summed["#work"]
		if total != 5956 || !ok {
			t.Errorf("summed['#work'] = %v, want %v", total, 5956)
		}
	} else {
		t.Errorf("parsing error %v", error)
	}
}
func TestWithStopSumFile(t *testing.T) {
	content := []string{
		"Mon Oct 28 20:00:00 +0100 2013;#play;",
		"Mon Oct 28 21:45:33 +0100 2013;stop;",
		"Mon Oct 28 21:45:33 +0100 2013;#work;",
		"Mon Oct 28 23:24:49 +0100 2013;stop;",
	}
	var summed, error = SumFile(content)

	if error == nil {
		var total, ok = summed["#play"]
		if total != 6333 || !ok {
			t.Errorf("summed['#play'] = %v, want %v", total, 6333)
		}
		total, ok = summed["#work"]
		if total != 5956 || !ok {
			t.Errorf("summed['#work'] = %v, want %v", total, 5956)
		}
	} else {
		t.Errorf("parsing error %v", error)
	}
}
<commit_msg>update tests<commit_after>package main

import (
	"os"
	"testing"
	"time"
	"bytes"
	"io"
)

// change traq env to use our fixtures
func WithFakeEnv(block func()) {
  oldEnv := os.Getenv("TRAQ_DATA_DIR")
  path, _ := os.Getwd()
  os.Setenv("TRAQ_DATA_DIR", path + "/fixtures")

  block()

  os.Setenv("TRAQ_DATA_DIR", oldEnv)
}

// capture output written to os.Stdout and return it
func CaptureStdout(block func()) string {
  old := os.Stdout // keep backup of the real stdout
  r, w, _ := os.Pipe()
  os.Stdout = w

  block()

  outC := make(chan string)
  // copy the output in a separate goroutine so printing can't block indefinitely
  go func() {
      var buf bytes.Buffer
      io.Copy(&buf, r)
      outC <- buf.String()
  }()

  // back to normal state
  w.Close()
  os.Stdout = old

  return <-outC
}

func TestDatesInMonth(t *testing.T) {
  dates := DatesInMonth(1986, 9)

  if len(dates) != 30 {
    t.Errorf("expected 30 days in Sep 1986, got %v", len(dates))
  }

  if dates[0].Weekday() != time.Monday {
    t.Errorf("Started on a Monday, got %v", dates[0].Weekday())
  }

  if dates[len(dates) - 1].Weekday() != time.Tuesday {
    t.Errorf("Ended on a Tuesday, got %v", dates[len(dates) - 1].Weekday())
  }
}

func TestPrintDate(t *testing.T) {
  out := CaptureStdout(func() {
    WithFakeEnv(func() {
      PrintDate("example", time.Date(1986, 9, 3, 0, 0, 0, 0, time.UTC))
    })
  })

  expected :=
`Wed Sep 03 20:00:00 +0100 1986;#birth;comment
Wed Sep 03 21:45:33 +0100 1986;#chillout;
Wed Sep 03 23:24:49 +0100 1986;stop;
%%
`
  if out != expected {
    t.Errorf("unexpected PrintDate output. Expected '%v' got '%v'", expected, out)
  }
}

func TestEvaluateDate(t *testing.T) {
  out := CaptureStdout(func() {
    WithFakeEnv(func() {
      EvaluateDate(ContentLoader, "example", time.Date(1986, 9, 3, 0, 0, 0, 0, time.UTC))
    })
  })

  expected :=
`1986-09-03
#birth:1.7592
#chillout:1.6544
%%
`

  if out != expected {
    t.Errorf("unexpected EvaluateDate output. Expected '%v' got '%v'", expected, out)
  }
}

func TestEntry(t *testing.T) {
  expected := `Wed Sep 3 12:00:00 +0000 1986;#test;
`

  if entry := Entry(time.Date(1986, 9, 3, 12, 0, 0, 0, time.UTC), "#test"); entry != expected {
    t.Errorf("got wrong entry. Expected '%v' got '%v'", expected, entry)
  }
}

func TestFilePath(t *testing.T) {
	var path string = FilePath("example", time.Date(1986, 9, 3, 0, 0, 0, 0, time.UTC))

	if path != os.Getenv("TRAQ_DATA_DIR")+"/example/1986/1986-09-03" {
		t.Errorf("FilePath = %v, want %v", path, os.Getenv("TRAQ_DATA_DIR")+"/example/1986/1986-09-03")
	}
}

func TestEmptySumFile(t *testing.T) {
	content := []string{""}
	var summed, error = SumFile(content)

	if error == nil {
		var total, ok = summed["#work"]
		if ok {
			t.Errorf("summed['#work'] = %v, should not exist", total)
		}
	} else {
		t.Errorf("parsing error %v", error)
	}
}

func TestSimpleSumFile(t *testing.T) {
	content := []string{
		"Mon Oct 28 21:45:33 +0100 2013;#work;",
		"Mon Oct 28 23:24:49 +0100 2013;stop;",
	}
	var summed, error = SumFile(content)

	if error == nil {
		var total, ok = summed["#work"]
		if total != 5956 || !ok {
			t.Errorf("summed['#work'] = %v, want %v", total, 5956)
		}
	} else {
		t.Errorf("parsing error %v", error)
	}
}

func TestNoStopSumFile(t *testing.T) {
	content := []string{
		"Mon Oct 28 20:00:00 +0100 2013;#play;",
		"Mon Oct 28 21:45:33 +0100 2013;#work;",
		"Mon Oct 28 23:24:49 +0100 2013;stop;",
	}
	var summed, error = SumFile(content)

	if error == nil {
		var total, ok = summed["#play"]
		if total != 6333 || !ok {
			t.Errorf("summed['#play'] = %v, want %v", total, 6333)
		}
		total, ok = summed["#work"]
		if total != 5956 || !ok {
			t.Errorf("summed['#work'] = %v, want %v", total, 5956)
		}
	} else {
		t.Errorf("parsing error %v", error)
	}
}
func TestWithStopSumFile(t *testing.T) {
	content := []string{
		"Mon Oct 28 20:00:00 +0100 2013;#play;",
		"Mon Oct 28 21:45:33 +0100 2013;stop;",
		"Mon Oct 28 21:45:33 +0100 2013;#work;",
		"Mon Oct 28 23:24:49 +0100 2013;stop;",
	}
	var summed, error = SumFile(content)

	if error == nil {
		var total, ok = summed["#play"]
		if total != 6333 || !ok {
			t.Errorf("summed['#play'] = %v, want %v", total, 6333)
		}
		total, ok = summed["#work"]
		if total != 5956 || !ok {
			t.Errorf("summed['#work'] = %v, want %v", total, 5956)
		}
	} else {
		t.Errorf("parsing error %v", error)
	}
}
<|endoftext|>