<commit_before>// Copyright 2015 The Prometheus 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 discovery

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

	"github.com/prometheus/log"
	"github.com/samuel/go-zookeeper/zk"

	clientmodel "github.com/prometheus/client_golang/model"

	"github.com/prometheus/prometheus/config"
)

const (
	serversetSourcePrefix = "serverset"

	serversetNodePrefix = "member_"

	serversetLabelPrefix         = clientmodel.MetaLabelPrefix + "serverset_"
	serversetStatusLabel         = serversetLabelPrefix + "status"
	serversetPathLabel           = serversetLabelPrefix + "path"
	serversetEndpointLabelPrefix = serversetLabelPrefix + "endpoint"
)

var (
	invalidLabelCharRE = regexp.MustCompile(`[^a-zA-Z0-9_]`)
)

type serversetMember struct {
	ServiceEndpoint     serversetEndpoint
	AdditionalEndpoints map[string]serversetEndpoint
	Status              string `json:"status"`
}

type serversetEndpoint struct {
	Host string
	Port int
}

type ZookeeperLogger struct {
}

// Implements zk.Logger
func (zl ZookeeperLogger) Printf(s string, i ...interface{}) {
	log.Infof(s, i...)
}

// ServersetDiscovery retrieves target information from a Serverset server
// and updates them via watches.
type ServersetDiscovery struct {
	conf      *config.ServersetSDConfig
	conn      *zk.Conn
	mu        sync.RWMutex
	sources   map[string]*config.TargetGroup
	sdUpdates *chan<- *config.TargetGroup
	updates   chan zookeeperTreeCacheEvent
	runDone   chan struct{}
	treeCache *zookeeperTreeCache
}

// NewServersetDiscovery returns a new ServersetDiscovery for the given config.
func NewServersetDiscovery(conf *config.ServersetSDConfig) *ServersetDiscovery {
	conn, _, err := zk.Connect(conf.Servers, time.Duration(conf.Timeout))
	conn.SetLogger(ZookeeperLogger{})
	if err != nil {
		return nil
	}
	updates := make(chan zookeeperTreeCacheEvent)
	sd := &ServersetDiscovery{
		conf:    conf,
		conn:    conn,
		updates: updates,
		sources: map[string]*config.TargetGroup{},
		runDone: make(chan struct{}),
	}
	go sd.processUpdates()
	sd.treeCache = NewZookeeperTreeCache(conn, conf.Paths[0], updates)
	return sd
}

// Sources implements the TargetProvider interface.
func (sd *ServersetDiscovery) Sources() []string {
	sd.mu.RLock()
	defer sd.mu.RUnlock()
	srcs := []string{}
	for t := range sd.sources {
		srcs = append(srcs, t)
	}
	return srcs
}

func (sd *ServersetDiscovery) processUpdates() {
	defer sd.conn.Close()
	for event := range sd.updates {
		tg := &config.TargetGroup{
			Source: serversetSourcePrefix + event.Path,
		}
		sd.mu.Lock()
		if event.Data != nil {
			labelSet, err := parseServersetMember(*event.Data, event.Path)
			if err == nil {
				tg.Targets = []clientmodel.LabelSet{*labelSet}
				sd.sources[event.Path] = tg
			} else {
				delete(sd.sources, event.Path)
			}
		} else {
			delete(sd.sources, event.Path)
		}
		sd.mu.Unlock()
		if sd.sdUpdates != nil {
			*sd.sdUpdates <- tg
		}
	}

	if sd.sdUpdates != nil {
		close(*sd.sdUpdates)
	}
}

// Run implements the TargetProvider interface.
func (sd *ServersetDiscovery) Run(ch chan<- *config.TargetGroup) {
	// Send on everything we have seen so far.
	sd.mu.Lock()
	for _, targetGroup := range sd.sources {
		ch <- targetGroup
	}
	// Tell processUpdates to send future updates.
	sd.sdUpdates = &ch
	sd.mu.Unlock()

	<-sd.runDone
	sd.treeCache.Stop()
}

// Stop implements the TargetProvider interface.
func (sd *ServersetDiscovery) Stop() {
	log.Debugf("Stopping serverset service discovery for %s %s", sd.conf.Servers, sd.conf.Paths)

	// Terminate Run.
	sd.runDone <- struct{}{}

	log.Debugf("Serverset service discovery for %s %s stopped", sd.conf.Servers, sd.conf.Paths)
}

func parseServersetMember(data []byte, path string) (*clientmodel.LabelSet, error) {
	member := serversetMember{}
	err := json.Unmarshal(data, &member)
	if err != nil {
		return nil, fmt.Errorf("error unmarshaling serverset member %q: %s", path, err)
	}

	labels := clientmodel.LabelSet{}
	labels[serversetPathLabel] = clientmodel.LabelValue(path)
	labels[clientmodel.AddressLabel] = clientmodel.LabelValue(
		fmt.Sprintf("%s:%d", member.ServiceEndpoint.Host, member.ServiceEndpoint.Port))

	labels[serversetEndpointLabelPrefix+"_host"] = clientmodel.LabelValue(member.ServiceEndpoint.Host)
	labels[serversetEndpointLabelPrefix+"_port"] = clientmodel.LabelValue(fmt.Sprintf("%d", member.ServiceEndpoint.Port))

	for name, endpoint := range member.AdditionalEndpoints {
		cleanName := clientmodel.LabelName(invalidLabelCharRE.ReplaceAllString(name, "_"))
		labels[serversetEndpointLabelPrefix+"_host_"+cleanName] = clientmodel.LabelValue(
			endpoint.Host)
		labels[serversetEndpointLabelPrefix+"_port_"+cleanName] = clientmodel.LabelValue(
			fmt.Sprintf("%d", endpoint.Port))

	}

	labels[serversetStatusLabel] = clientmodel.LabelValue(member.Status)

	return &labels, nil
}

type zookeeperTreeCache struct {
	conn     *zk.Conn
	prefix   string
	events   chan zookeeperTreeCacheEvent
	zkEvents chan zk.Event
	stop     chan struct{}
	head     *zookeeperTreeCacheNode
}

type zookeeperTreeCacheEvent struct {
	Path string
	Data *[]byte
}

type zookeeperTreeCacheNode struct {
	data     *[]byte
	events   chan zk.Event
	done     chan struct{}
	stopped  bool
	children map[string]*zookeeperTreeCacheNode
}

func NewZookeeperTreeCache(conn *zk.Conn, path string, events chan zookeeperTreeCacheEvent) *zookeeperTreeCache {
	tc := &zookeeperTreeCache{
		conn:   conn,
		prefix: path,
		events: events,
		stop:   make(chan struct{}),
	}
	tc.head = &zookeeperTreeCacheNode{
		events:   make(chan zk.Event),
		children: map[string]*zookeeperTreeCacheNode{},
	}
	err := tc.recursiveNodeUpdate(path, tc.head)
	if err != nil {
		log.Errorf("Error during initial read of Zookeeper: %s", err)
	}
	go tc.loop(err != nil)
	return tc
}

func (tc *zookeeperTreeCache) Stop() {
	tc.stop <- struct{}{}
}

func (tc *zookeeperTreeCache) loop(failureMode bool) {
	retryChan := make(chan struct{})

	failure := func() {
		failureMode = true
		time.AfterFunc(time.Second*10, func() {
			retryChan <- struct{}{}
		})
	}
	if failureMode {
		failure()
	}

	for {
		select {
		case ev := <-tc.head.events:
			log.Debugf("Received Zookeeper event: %s", ev)
			if failureMode {
				continue
			}
			if ev.Type == zk.EventNotWatching {
				log.Infof("Lost connection to Zookeeper.")
				failure()
			} else {
				path := strings.TrimPrefix(ev.Path, tc.prefix)
				parts := strings.Split(path, "/")
				node := tc.head
				for _, part := range parts[1:] {
					childNode := node.children[part]
					if childNode == nil {
						childNode = &zookeeperTreeCacheNode{
							events:   tc.head.events,
							children: map[string]*zookeeperTreeCacheNode{},
							done:     make(chan struct{}, 1),
						}
						node.children[part] = childNode
					}
					node = childNode
				}
				err := tc.recursiveNodeUpdate(ev.Path, node)
				if err != nil {
					log.Errorf("Error during processing of Zookeeper event: %s", err)
					failure()
				}
			}
		case <-retryChan:
			log.Infof("Attempting to resync state with Zookeeper")
			err := tc.recursiveNodeUpdate(tc.prefix, tc.head)
			if err == nil {
				failureMode = false
			} else {
				log.Errorf("Error during Zookeeper resync: %s", err)
				failure()
			}
		case <-tc.stop:
			close(tc.events)
			return
		}
	}
}

func (tc *zookeeperTreeCache) recursiveNodeUpdate(path string, node *zookeeperTreeCacheNode) error {
	data, _, dataWatcher, err := tc.conn.GetW(path)
	if err == zk.ErrNoNode {
		tc.recursiveDelete(path, node)
		return nil
	} else if err != nil {
		return err
	}

	if node.data == nil || !bytes.Equal(*node.data, data) {
		node.data = &data
		tc.events <- zookeeperTreeCacheEvent{Path: path, Data: node.data}
	}

	children, _, childWatcher, err := tc.conn.ChildrenW(path)
	if err == zk.ErrNoNode {
		tc.recursiveDelete(path, node)
		return nil
	} else if err != nil {
		return err
	}

	currentChildren := map[string]struct{}{}
	for _, child := range children {
		currentChildren[child] = struct{}{}
		childNode := node.children[child]
		// Does not already exists, create it.
		if childNode == nil {
			node.children[child] = &zookeeperTreeCacheNode{
				events:   node.events,
				children: map[string]*zookeeperTreeCacheNode{},
				done:     make(chan struct{}, 1),
			}
		}
		err = tc.recursiveNodeUpdate(path+"/"+child, node.children[child])
		if err != nil {
			return err
		}
	}
	// Remove nodes that no longer exist
	for name, childNode := range node.children {
		if _, ok := currentChildren[name]; !ok || node.data == nil {
			tc.recursiveDelete(path+"/"+name, childNode)
			delete(node.children, name)
		}
	}

	go func() {
		// Pass up zookeeper events, until the node is deleted.
		select {
		case event := <-dataWatcher:
			node.events <- event
		case event := <-childWatcher:
			node.events <- event
		case <-node.done:
		}
	}()
	return nil
}

func (tc *zookeeperTreeCache) recursiveDelete(path string, node *zookeeperTreeCacheNode) {
	if !node.stopped {
		node.done <- struct{}{}
		node.stopped = true
	}
	if node.data != nil {
		tc.events <- zookeeperTreeCacheEvent{Path: path, Data: nil}
		node.data = nil
	}
	for name, childNode := range node.children {
		tc.recursiveDelete(path+"/"+name, childNode)
	}
}
<commit_msg>retrieval: Handle serverset node not existing.<commit_after>// Copyright 2015 The Prometheus 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 discovery

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

	"github.com/prometheus/log"
	"github.com/samuel/go-zookeeper/zk"

	clientmodel "github.com/prometheus/client_golang/model"

	"github.com/prometheus/prometheus/config"
)

const (
	serversetSourcePrefix = "serverset"

	serversetNodePrefix = "member_"

	serversetLabelPrefix         = clientmodel.MetaLabelPrefix + "serverset_"
	serversetStatusLabel         = serversetLabelPrefix + "status"
	serversetPathLabel           = serversetLabelPrefix + "path"
	serversetEndpointLabelPrefix = serversetLabelPrefix + "endpoint"
)

var (
	invalidLabelCharRE = regexp.MustCompile(`[^a-zA-Z0-9_]`)
)

type serversetMember struct {
	ServiceEndpoint     serversetEndpoint
	AdditionalEndpoints map[string]serversetEndpoint
	Status              string `json:"status"`
}

type serversetEndpoint struct {
	Host string
	Port int
}

type ZookeeperLogger struct {
}

// Implements zk.Logger
func (zl ZookeeperLogger) Printf(s string, i ...interface{}) {
	log.Infof(s, i...)
}

// ServersetDiscovery retrieves target information from a Serverset server
// and updates them via watches.
type ServersetDiscovery struct {
	conf      *config.ServersetSDConfig
	conn      *zk.Conn
	mu        sync.RWMutex
	sources   map[string]*config.TargetGroup
	sdUpdates *chan<- *config.TargetGroup
	updates   chan zookeeperTreeCacheEvent
	runDone   chan struct{}
	treeCache *zookeeperTreeCache
}

// NewServersetDiscovery returns a new ServersetDiscovery for the given config.
func NewServersetDiscovery(conf *config.ServersetSDConfig) *ServersetDiscovery {
	conn, _, err := zk.Connect(conf.Servers, time.Duration(conf.Timeout))
	conn.SetLogger(ZookeeperLogger{})
	if err != nil {
		return nil
	}
	updates := make(chan zookeeperTreeCacheEvent)
	sd := &ServersetDiscovery{
		conf:    conf,
		conn:    conn,
		updates: updates,
		sources: map[string]*config.TargetGroup{},
		runDone: make(chan struct{}),
	}
	go sd.processUpdates()
	sd.treeCache = NewZookeeperTreeCache(conn, conf.Paths[0], updates)
	return sd
}

// Sources implements the TargetProvider interface.
func (sd *ServersetDiscovery) Sources() []string {
	sd.mu.RLock()
	defer sd.mu.RUnlock()
	srcs := []string{}
	for t := range sd.sources {
		srcs = append(srcs, t)
	}
	return srcs
}

func (sd *ServersetDiscovery) processUpdates() {
	defer sd.conn.Close()
	for event := range sd.updates {
		tg := &config.TargetGroup{
			Source: serversetSourcePrefix + event.Path,
		}
		sd.mu.Lock()
		if event.Data != nil {
			labelSet, err := parseServersetMember(*event.Data, event.Path)
			if err == nil {
				tg.Targets = []clientmodel.LabelSet{*labelSet}
				sd.sources[event.Path] = tg
			} else {
				delete(sd.sources, event.Path)
			}
		} else {
			delete(sd.sources, event.Path)
		}
		sd.mu.Unlock()
		if sd.sdUpdates != nil {
			*sd.sdUpdates <- tg
		}
	}

	if sd.sdUpdates != nil {
		close(*sd.sdUpdates)
	}
}

// Run implements the TargetProvider interface.
func (sd *ServersetDiscovery) Run(ch chan<- *config.TargetGroup) {
	// Send on everything we have seen so far.
	sd.mu.Lock()
	for _, targetGroup := range sd.sources {
		ch <- targetGroup
	}
	// Tell processUpdates to send future updates.
	sd.sdUpdates = &ch
	sd.mu.Unlock()

	<-sd.runDone
	sd.treeCache.Stop()
}

// Stop implements the TargetProvider interface.
func (sd *ServersetDiscovery) Stop() {
	log.Debugf("Stopping serverset service discovery for %s %s", sd.conf.Servers, sd.conf.Paths)

	// Terminate Run.
	sd.runDone <- struct{}{}

	log.Debugf("Serverset service discovery for %s %s stopped", sd.conf.Servers, sd.conf.Paths)
}

func parseServersetMember(data []byte, path string) (*clientmodel.LabelSet, error) {
	member := serversetMember{}
	err := json.Unmarshal(data, &member)
	if err != nil {
		return nil, fmt.Errorf("error unmarshaling serverset member %q: %s", path, err)
	}

	labels := clientmodel.LabelSet{}
	labels[serversetPathLabel] = clientmodel.LabelValue(path)
	labels[clientmodel.AddressLabel] = clientmodel.LabelValue(
		fmt.Sprintf("%s:%d", member.ServiceEndpoint.Host, member.ServiceEndpoint.Port))

	labels[serversetEndpointLabelPrefix+"_host"] = clientmodel.LabelValue(member.ServiceEndpoint.Host)
	labels[serversetEndpointLabelPrefix+"_port"] = clientmodel.LabelValue(fmt.Sprintf("%d", member.ServiceEndpoint.Port))

	for name, endpoint := range member.AdditionalEndpoints {
		cleanName := clientmodel.LabelName(invalidLabelCharRE.ReplaceAllString(name, "_"))
		labels[serversetEndpointLabelPrefix+"_host_"+cleanName] = clientmodel.LabelValue(
			endpoint.Host)
		labels[serversetEndpointLabelPrefix+"_port_"+cleanName] = clientmodel.LabelValue(
			fmt.Sprintf("%d", endpoint.Port))

	}

	labels[serversetStatusLabel] = clientmodel.LabelValue(member.Status)

	return &labels, nil
}

type zookeeperTreeCache struct {
	conn     *zk.Conn
	prefix   string
	events   chan zookeeperTreeCacheEvent
	zkEvents chan zk.Event
	stop     chan struct{}
	head     *zookeeperTreeCacheNode
}

type zookeeperTreeCacheEvent struct {
	Path string
	Data *[]byte
}

type zookeeperTreeCacheNode struct {
	data     *[]byte
	events   chan zk.Event
	done     chan struct{}
	stopped  bool
	children map[string]*zookeeperTreeCacheNode
}

func NewZookeeperTreeCache(conn *zk.Conn, path string, events chan zookeeperTreeCacheEvent) *zookeeperTreeCache {
	tc := &zookeeperTreeCache{
		conn:   conn,
		prefix: path,
		events: events,
		stop:   make(chan struct{}),
	}
	tc.head = &zookeeperTreeCacheNode{
		events:   make(chan zk.Event),
		children: map[string]*zookeeperTreeCacheNode{},
		stopped:  true,
	}
	err := tc.recursiveNodeUpdate(path, tc.head)
	if err != nil {
		log.Errorf("Error during initial read of Zookeeper: %s", err)
	}
	go tc.loop(err != nil)
	return tc
}

func (tc *zookeeperTreeCache) Stop() {
	tc.stop <- struct{}{}
}

func (tc *zookeeperTreeCache) loop(failureMode bool) {
	retryChan := make(chan struct{})

	failure := func() {
		failureMode = true
		time.AfterFunc(time.Second*10, func() {
			retryChan <- struct{}{}
		})
	}
	if failureMode {
		failure()
	}

	for {
		select {
		case ev := <-tc.head.events:
			log.Debugf("Received Zookeeper event: %s", ev)
			if failureMode {
				continue
			}
			if ev.Type == zk.EventNotWatching {
				log.Infof("Lost connection to Zookeeper.")
				failure()
			} else {
				path := strings.TrimPrefix(ev.Path, tc.prefix)
				parts := strings.Split(path, "/")
				node := tc.head
				for _, part := range parts[1:] {
					childNode := node.children[part]
					if childNode == nil {
						childNode = &zookeeperTreeCacheNode{
							events:   tc.head.events,
							children: map[string]*zookeeperTreeCacheNode{},
							done:     make(chan struct{}, 1),
						}
						node.children[part] = childNode
					}
					node = childNode
				}
				err := tc.recursiveNodeUpdate(ev.Path, node)
				if err != nil {
					log.Errorf("Error during processing of Zookeeper event: %s", err)
					failure()
				} else if tc.head.data == nil {
					log.Errorf("Error during processing of Zookeeper event: path %s no longer exists", tc.prefix)
					failure()
				}
			}
		case <-retryChan:
			log.Infof("Attempting to resync state with Zookeeper")
			err := tc.recursiveNodeUpdate(tc.prefix, tc.head)
			if err != nil {
				log.Errorf("Error during Zookeeper resync: %s", err)
				failure()
			} else {
				log.Infof("Zookeeper resync successful")
				failureMode = false
			}
		case <-tc.stop:
			close(tc.events)
			return
		}
	}
}

func (tc *zookeeperTreeCache) recursiveNodeUpdate(path string, node *zookeeperTreeCacheNode) error {
	data, _, dataWatcher, err := tc.conn.GetW(path)
	if err == zk.ErrNoNode {
		tc.recursiveDelete(path, node)
		if node == tc.head {
			return fmt.Errorf("path %s does not exist", path)
		}
		return nil
	} else if err != nil {
		return err
	}

	if node.data == nil || !bytes.Equal(*node.data, data) {
		node.data = &data
		tc.events <- zookeeperTreeCacheEvent{Path: path, Data: node.data}
	}

	children, _, childWatcher, err := tc.conn.ChildrenW(path)
	if err == zk.ErrNoNode {
		tc.recursiveDelete(path, node)
		return nil
	} else if err != nil {
		return err
	}

	currentChildren := map[string]struct{}{}
	for _, child := range children {
		currentChildren[child] = struct{}{}
		childNode := node.children[child]
		// Does not already exists, create it.
		if childNode == nil {
			node.children[child] = &zookeeperTreeCacheNode{
				events:   node.events,
				children: map[string]*zookeeperTreeCacheNode{},
				done:     make(chan struct{}, 1),
			}
		}
		err = tc.recursiveNodeUpdate(path+"/"+child, node.children[child])
		if err != nil {
			return err
		}
	}
	// Remove nodes that no longer exist
	for name, childNode := range node.children {
		if _, ok := currentChildren[name]; !ok || node.data == nil {
			tc.recursiveDelete(path+"/"+name, childNode)
			delete(node.children, name)
		}
	}

	go func() {
		// Pass up zookeeper events, until the node is deleted.
		select {
		case event := <-dataWatcher:
			node.events <- event
		case event := <-childWatcher:
			node.events <- event
		case <-node.done:
		}
	}()
	return nil
}

func (tc *zookeeperTreeCache) recursiveDelete(path string, node *zookeeperTreeCacheNode) {
	if !node.stopped {
		node.done <- struct{}{}
		node.stopped = true
	}
	if node.data != nil {
		tc.events <- zookeeperTreeCacheEvent{Path: path, Data: nil}
		node.data = nil
	}
	for name, childNode := range node.children {
		tc.recursiveDelete(path+"/"+name, childNode)
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package user

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

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/base"
	"code.gitea.io/gitea/modules/context"
	"code.gitea.io/gitea/modules/setting"
	"code.gitea.io/gitea/modules/structs"
)

const (
	tplNotification    base.TplName = "user/notification/notification"
	tplNotificationDiv base.TplName = "user/notification/notification_div"
)

// GetNotificationCount is the middleware that sets the notification count in the context
func GetNotificationCount(c *context.Context) {
	if strings.HasPrefix(c.Req.URL.Path, "/api") {
		return
	}

	if !c.IsSigned {
		return
	}

	c.Data["NotificationUnreadCount"] = func() int64 {
		count, err := models.GetNotificationCount(c, c.Doer, models.NotificationStatusUnread)
		if err != nil {
			c.ServerError("GetNotificationCount", err)
			return -1
		}

		return count
	}
}

// Notifications is the notifications page
func Notifications(c *context.Context) {
	getNotifications(c)
	if c.Written() {
		return
	}
	if c.FormBool("div-only") {
		c.Data["SequenceNumber"] = c.FormString("sequence-number")
		c.HTML(http.StatusOK, tplNotificationDiv)
		return
	}
	c.HTML(http.StatusOK, tplNotification)
}

func getNotifications(c *context.Context) {
	var (
		keyword = c.FormTrim("q")
		status  models.NotificationStatus
		page    = c.FormInt("page")
		perPage = c.FormInt("perPage")
	)
	if page < 1 {
		page = 1
	}
	if perPage < 1 {
		perPage = 20
	}

	switch keyword {
	case "read":
		status = models.NotificationStatusRead
	default:
		status = models.NotificationStatusUnread
	}

	total, err := models.GetNotificationCount(c, c.Doer, status)
	if err != nil {
		c.ServerError("ErrGetNotificationCount", err)
		return
	}

	// redirect to last page if request page is more than total pages
	pager := context.NewPagination(int(total), perPage, page, 5)
	if pager.Paginater.Current() < page {
		c.Redirect(fmt.Sprintf("%s/notifications?q=%s&page=%d", setting.AppSubURL, url.QueryEscape(c.FormString("q")), pager.Paginater.Current()))
		return
	}

	statuses := []models.NotificationStatus{status, models.NotificationStatusPinned}
	notifications, err := models.NotificationsForUser(c, c.Doer, statuses, page, perPage)
	if err != nil {
		c.ServerError("ErrNotificationsForUser", err)
		return
	}

	failCount := 0

	repos, failures, err := notifications.LoadRepos()
	if err != nil {
		c.ServerError("LoadRepos", err)
		return
	}
	notifications = notifications.Without(failures)
	if err := repos.LoadAttributes(); err != nil {
		c.ServerError("LoadAttributes", err)
		return
	}
	failCount += len(failures)

	failures, err = notifications.LoadIssues()
	if err != nil {
		c.ServerError("LoadIssues", err)
		return
	}
	notifications = notifications.Without(failures)
	failCount += len(failures)

	failures, err = notifications.LoadComments()
	if err != nil {
		c.ServerError("LoadComments", err)
		return
	}
	notifications = notifications.Without(failures)
	failCount += len(failures)

	if failCount > 0 {
		c.Flash.Error(fmt.Sprintf("ERROR: %d notifications were removed due to missing parts - check the logs", failCount))
	}

	c.Data["Title"] = c.Tr("notifications")
	c.Data["Keyword"] = keyword
	c.Data["Status"] = status
	c.Data["Notifications"] = notifications

	pager.SetDefaultParams(c)
	c.Data["Page"] = pager
}

// NotificationStatusPost is a route for changing the status of a notification
func NotificationStatusPost(c *context.Context) {
	var (
		notificationID = c.FormInt64("notification_id")
		statusStr      = c.FormString("status")
		status         models.NotificationStatus
	)

	switch statusStr {
	case "read":
		status = models.NotificationStatusRead
	case "unread":
		status = models.NotificationStatusUnread
	case "pinned":
		status = models.NotificationStatusPinned
	default:
		c.ServerError("InvalidNotificationStatus", errors.New("Invalid notification status"))
		return
	}

	if _, err := models.SetNotificationStatus(notificationID, c.Doer, status); err != nil {
		c.ServerError("SetNotificationStatus", err)
		return
	}

	if !c.FormBool("noredirect") {
		url := fmt.Sprintf("%s/notifications?page=%s", setting.AppSubURL, url.QueryEscape(c.FormString("page")))
		c.Redirect(url, http.StatusSeeOther)
	}

	getNotifications(c)
	if c.Written() {
		return
	}
	c.Data["Link"] = setting.AppURL + "notifications"
	c.Data["SequenceNumber"] = c.Req.PostFormValue("sequence-number")

	c.HTML(http.StatusOK, tplNotificationDiv)
}

// NotificationPurgePost is a route for 'purging' the list of notifications - marking all unread as read
func NotificationPurgePost(c *context.Context) {
	err := models.UpdateNotificationStatuses(c.Doer, models.NotificationStatusUnread, models.NotificationStatusRead)
	if err != nil {
		c.ServerError("ErrUpdateNotificationStatuses", err)
		return
	}

	c.Redirect(setting.AppSubURL+"/notifications", http.StatusSeeOther)
}

// NewAvailable returns the notification counts
func NewAvailable(ctx *context.Context) {
	ctx.JSON(http.StatusOK, structs.NotificationCount{New: models.CountUnread(ctx, ctx.Doer.ID)})
}
<commit_msg>Prevent loop if there is an error in GetNotificationCount (#19799)<commit_after>// Copyright 2019 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package user

import (
	goctx "context"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/base"
	"code.gitea.io/gitea/modules/context"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"
	"code.gitea.io/gitea/modules/structs"
)

const (
	tplNotification    base.TplName = "user/notification/notification"
	tplNotificationDiv base.TplName = "user/notification/notification_div"
)

// GetNotificationCount is the middleware that sets the notification count in the context
func GetNotificationCount(c *context.Context) {
	if strings.HasPrefix(c.Req.URL.Path, "/api") {
		return
	}

	if !c.IsSigned {
		return
	}

	c.Data["NotificationUnreadCount"] = func() int64 {
		count, err := models.GetNotificationCount(c, c.Doer, models.NotificationStatusUnread)
		if err != nil {
			if err != goctx.Canceled {
				log.Error("Unable to GetNotificationCount for user:%-v: %v", c.Doer, err)
			}
			return -1
		}

		return count
	}
}

// Notifications is the notifications page
func Notifications(c *context.Context) {
	getNotifications(c)
	if c.Written() {
		return
	}
	if c.FormBool("div-only") {
		c.Data["SequenceNumber"] = c.FormString("sequence-number")
		c.HTML(http.StatusOK, tplNotificationDiv)
		return
	}
	c.HTML(http.StatusOK, tplNotification)
}

func getNotifications(c *context.Context) {
	var (
		keyword = c.FormTrim("q")
		status  models.NotificationStatus
		page    = c.FormInt("page")
		perPage = c.FormInt("perPage")
	)
	if page < 1 {
		page = 1
	}
	if perPage < 1 {
		perPage = 20
	}

	switch keyword {
	case "read":
		status = models.NotificationStatusRead
	default:
		status = models.NotificationStatusUnread
	}

	total, err := models.GetNotificationCount(c, c.Doer, status)
	if err != nil {
		c.ServerError("ErrGetNotificationCount", err)
		return
	}

	// redirect to last page if request page is more than total pages
	pager := context.NewPagination(int(total), perPage, page, 5)
	if pager.Paginater.Current() < page {
		c.Redirect(fmt.Sprintf("%s/notifications?q=%s&page=%d", setting.AppSubURL, url.QueryEscape(c.FormString("q")), pager.Paginater.Current()))
		return
	}

	statuses := []models.NotificationStatus{status, models.NotificationStatusPinned}
	notifications, err := models.NotificationsForUser(c, c.Doer, statuses, page, perPage)
	if err != nil {
		c.ServerError("ErrNotificationsForUser", err)
		return
	}

	failCount := 0

	repos, failures, err := notifications.LoadRepos()
	if err != nil {
		c.ServerError("LoadRepos", err)
		return
	}
	notifications = notifications.Without(failures)
	if err := repos.LoadAttributes(); err != nil {
		c.ServerError("LoadAttributes", err)
		return
	}
	failCount += len(failures)

	failures, err = notifications.LoadIssues()
	if err != nil {
		c.ServerError("LoadIssues", err)
		return
	}
	notifications = notifications.Without(failures)
	failCount += len(failures)

	failures, err = notifications.LoadComments()
	if err != nil {
		c.ServerError("LoadComments", err)
		return
	}
	notifications = notifications.Without(failures)
	failCount += len(failures)

	if failCount > 0 {
		c.Flash.Error(fmt.Sprintf("ERROR: %d notifications were removed due to missing parts - check the logs", failCount))
	}

	c.Data["Title"] = c.Tr("notifications")
	c.Data["Keyword"] = keyword
	c.Data["Status"] = status
	c.Data["Notifications"] = notifications

	pager.SetDefaultParams(c)
	c.Data["Page"] = pager
}

// NotificationStatusPost is a route for changing the status of a notification
func NotificationStatusPost(c *context.Context) {
	var (
		notificationID = c.FormInt64("notification_id")
		statusStr      = c.FormString("status")
		status         models.NotificationStatus
	)

	switch statusStr {
	case "read":
		status = models.NotificationStatusRead
	case "unread":
		status = models.NotificationStatusUnread
	case "pinned":
		status = models.NotificationStatusPinned
	default:
		c.ServerError("InvalidNotificationStatus", errors.New("Invalid notification status"))
		return
	}

	if _, err := models.SetNotificationStatus(notificationID, c.Doer, status); err != nil {
		c.ServerError("SetNotificationStatus", err)
		return
	}

	if !c.FormBool("noredirect") {
		url := fmt.Sprintf("%s/notifications?page=%s", setting.AppSubURL, url.QueryEscape(c.FormString("page")))
		c.Redirect(url, http.StatusSeeOther)
	}

	getNotifications(c)
	if c.Written() {
		return
	}
	c.Data["Link"] = setting.AppURL + "notifications"
	c.Data["SequenceNumber"] = c.Req.PostFormValue("sequence-number")

	c.HTML(http.StatusOK, tplNotificationDiv)
}

// NotificationPurgePost is a route for 'purging' the list of notifications - marking all unread as read
func NotificationPurgePost(c *context.Context) {
	err := models.UpdateNotificationStatuses(c.Doer, models.NotificationStatusUnread, models.NotificationStatusRead)
	if err != nil {
		c.ServerError("ErrUpdateNotificationStatuses", err)
		return
	}

	c.Redirect(setting.AppSubURL+"/notifications", http.StatusSeeOther)
}

// NewAvailable returns the notification counts
func NewAvailable(ctx *context.Context) {
	ctx.JSON(http.StatusOK, structs.NotificationCount{New: models.CountUnread(ctx, ctx.Doer.ID)})
}
<|endoftext|>


<commit_before>package mdbd

import (
	"bufio"
	"encoding/gob"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path"
	"reflect"
	"sort"
	"time"

	"github.com/Symantec/Dominator/lib/fsutil"
	jsonwriter "github.com/Symantec/Dominator/lib/json"
	"github.com/Symantec/Dominator/lib/log"
	"github.com/Symantec/Dominator/lib/mdb"
	"github.com/Symantec/Dominator/lib/srpc"
	"github.com/Symantec/Dominator/proto/mdbserver"
)

func startMdbDaemon(mdbFileName string, logger log.Logger) <-chan *mdb.Mdb {
	mdbChannel := make(chan *mdb.Mdb, 1)
	if *mdbServerHostname != "" && *mdbServerPortNum > 0 {
		go serverWatchDaemon(*mdbServerHostname, *mdbServerPortNum, mdbFileName,
			mdbChannel, logger)
	} else {
		go fileWatchDaemon(mdbFileName, mdbChannel, logger)
	}
	return mdbChannel
}

type genericDecoder interface {
	Decode(v interface{}) error
}

func fileWatchDaemon(mdbFileName string, mdbChannel chan<- *mdb.Mdb,
	logger log.Logger) {
	var lastMdb *mdb.Mdb
	for readCloser := range fsutil.WatchFile(mdbFileName, logger) {
		mdb := loadFile(readCloser, mdbFileName, logger)
		readCloser.Close()
		if mdb == nil {
			continue
		}
		compareStartTime := time.Now()
		if lastMdb == nil || !reflect.DeepEqual(lastMdb, mdb) {
			if lastMdb != nil {
				mdbCompareTimeDistribution.Add(time.Since(compareStartTime))
			}
			mdbChannel <- mdb
			lastMdb = mdb
		}
	}
}

func serverWatchDaemon(mdbServerHostname string, mdbServerPortNum uint,
	mdbFileName string, mdbChannel chan<- *mdb.Mdb, logger log.Logger) {
	if file, err := os.Open(mdbFileName); err == nil {
		fileMdb := loadFile(file, mdbFileName, logger)
		file.Close()
		if fileMdb != nil {
			sort.Sort(fileMdb)
			mdbChannel <- fileMdb
		}
	}
	address := fmt.Sprintf("%s:%d", mdbServerHostname, mdbServerPortNum)
	for ; ; time.Sleep(time.Second) {
		client, err := srpc.DialHTTP("tcp", address, time.Second*15)
		if err != nil {
			logger.Println(err)
			continue
		}
		conn, err := client.Call("MdbServer.GetMdbUpdates")
		if err != nil {
			logger.Println(err)
			client.Close()
			continue
		}
		decoder := gob.NewDecoder(conn)
		lastMdb := &mdb.Mdb{}
		for {
			var mdbUpdate mdbserver.MdbUpdate
			if err := decoder.Decode(&mdbUpdate); err != nil {
				logger.Println(err)
				break
			} else {
				lastMdb = processUpdate(lastMdb, mdbUpdate)
				sort.Sort(lastMdb)
				mdbChannel <- lastMdb
				if file, err := os.Create(mdbFileName + "~"); err != nil {
					logger.Println(err)
				} else {
					writer := bufio.NewWriter(file)
					var err error
					if isGob(mdbFileName) {
						encoder := gob.NewEncoder(writer)
						err = encoder.Encode(lastMdb.Machines)
					} else {
						err = jsonwriter.WriteWithIndent(writer, "    ",
							lastMdb.Machines)
					}
					if err != nil {
						logger.Println(err)
						os.Remove(mdbFileName + "~")
					} else {
						writer.Flush()
						file.Close()
						os.Rename(mdbFileName+"~", mdbFileName)
					}
				}
			}
		}
		conn.Close()
		client.Close()
	}
}

func loadFile(reader io.Reader, filename string, logger log.Logger) *mdb.Mdb {
	decoder := getDecoder(reader, filename)
	var mdb mdb.Mdb
	decodeStartTime := time.Now()
	if err := decoder.Decode(&mdb.Machines); err != nil {
		logger.Printf("Error decoding MDB data: %s\n", err)
		return nil
	}
	sortStartTime := time.Now()
	mdbDecodeTimeDistribution.Add(sortStartTime.Sub(decodeStartTime))
	sort.Sort(&mdb)
	mdbSortTimeDistribution.Add(time.Since(sortStartTime))
	return &mdb
}

func isGob(filename string) bool {
	switch path.Ext(filename) {
	case ".gob":
		return true
	default:
		return false
	}
}

func getDecoder(reader io.Reader, filename string) genericDecoder {
	if isGob(filename) {
		return gob.NewDecoder(reader)
	} else {
		return json.NewDecoder(reader)
	}
}

func processUpdate(oldMdb *mdb.Mdb, mdbUpdate mdbserver.MdbUpdate) *mdb.Mdb {
	newMdb := &mdb.Mdb{}
	if len(oldMdb.Machines) < 1 {
		newMdb.Machines = mdbUpdate.MachinesToAdd
		return newMdb
	}
	newMachines := make(map[string]mdb.Machine)
	for _, machine := range oldMdb.Machines {
		newMachines[machine.Hostname] = machine
	}
	for _, machine := range mdbUpdate.MachinesToAdd {
		newMachines[machine.Hostname] = machine
	}
	for _, machine := range mdbUpdate.MachinesToUpdate {
		newMachines[machine.Hostname] = machine
	}
	for _, name := range mdbUpdate.MachinesToDelete {
		delete(newMachines, name)
	}
	newMdb.Machines = make([]mdb.Machine, 0, len(newMachines))
	for _, machine := range newMachines {
		newMdb.Machines = append(newMdb.Machines, machine)
	}
	return newMdb
}
<commit_msg>Switch lib/mdb/mdbd to use connection Decode and Encode methods.<commit_after>package mdbd

import (
	"bufio"
	"encoding/gob"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path"
	"reflect"
	"sort"
	"time"

	"github.com/Symantec/Dominator/lib/fsutil"
	jsonwriter "github.com/Symantec/Dominator/lib/json"
	"github.com/Symantec/Dominator/lib/log"
	"github.com/Symantec/Dominator/lib/mdb"
	"github.com/Symantec/Dominator/lib/srpc"
	"github.com/Symantec/Dominator/proto/mdbserver"
)

func startMdbDaemon(mdbFileName string, logger log.Logger) <-chan *mdb.Mdb {
	mdbChannel := make(chan *mdb.Mdb, 1)
	if *mdbServerHostname != "" && *mdbServerPortNum > 0 {
		go serverWatchDaemon(*mdbServerHostname, *mdbServerPortNum, mdbFileName,
			mdbChannel, logger)
	} else {
		go fileWatchDaemon(mdbFileName, mdbChannel, logger)
	}
	return mdbChannel
}

type genericDecoder interface {
	Decode(v interface{}) error
}

func fileWatchDaemon(mdbFileName string, mdbChannel chan<- *mdb.Mdb,
	logger log.Logger) {
	var lastMdb *mdb.Mdb
	for readCloser := range fsutil.WatchFile(mdbFileName, logger) {
		mdb := loadFile(readCloser, mdbFileName, logger)
		readCloser.Close()
		if mdb == nil {
			continue
		}
		compareStartTime := time.Now()
		if lastMdb == nil || !reflect.DeepEqual(lastMdb, mdb) {
			if lastMdb != nil {
				mdbCompareTimeDistribution.Add(time.Since(compareStartTime))
			}
			mdbChannel <- mdb
			lastMdb = mdb
		}
	}
}

func serverWatchDaemon(mdbServerHostname string, mdbServerPortNum uint,
	mdbFileName string, mdbChannel chan<- *mdb.Mdb, logger log.Logger) {
	if file, err := os.Open(mdbFileName); err == nil {
		fileMdb := loadFile(file, mdbFileName, logger)
		file.Close()
		if fileMdb != nil {
			sort.Sort(fileMdb)
			mdbChannel <- fileMdb
		}
	}
	address := fmt.Sprintf("%s:%d", mdbServerHostname, mdbServerPortNum)
	for ; ; time.Sleep(time.Second) {
		client, err := srpc.DialHTTP("tcp", address, time.Second*15)
		if err != nil {
			logger.Println(err)
			continue
		}
		conn, err := client.Call("MdbServer.GetMdbUpdates")
		if err != nil {
			logger.Println(err)
			client.Close()
			continue
		}
		lastMdb := &mdb.Mdb{}
		for {
			var mdbUpdate mdbserver.MdbUpdate
			if err := conn.Decode(&mdbUpdate); err != nil {
				logger.Println(err)
				break
			} else {
				lastMdb = processUpdate(lastMdb, mdbUpdate)
				sort.Sort(lastMdb)
				mdbChannel <- lastMdb
				if file, err := os.Create(mdbFileName + "~"); err != nil {
					logger.Println(err)
				} else {
					writer := bufio.NewWriter(file)
					var err error
					if isGob(mdbFileName) {
						encoder := gob.NewEncoder(writer)
						err = encoder.Encode(lastMdb.Machines)
					} else {
						err = jsonwriter.WriteWithIndent(writer, "    ",
							lastMdb.Machines)
					}
					if err != nil {
						logger.Println(err)
						os.Remove(mdbFileName + "~")
					} else {
						writer.Flush()
						file.Close()
						os.Rename(mdbFileName+"~", mdbFileName)
					}
				}
			}
		}
		conn.Close()
		client.Close()
	}
}

func loadFile(reader io.Reader, filename string, logger log.Logger) *mdb.Mdb {
	decoder := getDecoder(reader, filename)
	var mdb mdb.Mdb
	decodeStartTime := time.Now()
	if err := decoder.Decode(&mdb.Machines); err != nil {
		logger.Printf("Error decoding MDB data: %s\n", err)
		return nil
	}
	sortStartTime := time.Now()
	mdbDecodeTimeDistribution.Add(sortStartTime.Sub(decodeStartTime))
	sort.Sort(&mdb)
	mdbSortTimeDistribution.Add(time.Since(sortStartTime))
	return &mdb
}

func isGob(filename string) bool {
	switch path.Ext(filename) {
	case ".gob":
		return true
	default:
		return false
	}
}

func getDecoder(reader io.Reader, filename string) genericDecoder {
	if isGob(filename) {
		return gob.NewDecoder(reader)
	} else {
		return json.NewDecoder(reader)
	}
}

func processUpdate(oldMdb *mdb.Mdb, mdbUpdate mdbserver.MdbUpdate) *mdb.Mdb {
	newMdb := &mdb.Mdb{}
	if len(oldMdb.Machines) < 1 {
		newMdb.Machines = mdbUpdate.MachinesToAdd
		return newMdb
	}
	newMachines := make(map[string]mdb.Machine)
	for _, machine := range oldMdb.Machines {
		newMachines[machine.Hostname] = machine
	}
	for _, machine := range mdbUpdate.MachinesToAdd {
		newMachines[machine.Hostname] = machine
	}
	for _, machine := range mdbUpdate.MachinesToUpdate {
		newMachines[machine.Hostname] = machine
	}
	for _, name := range mdbUpdate.MachinesToDelete {
		delete(newMachines, name)
	}
	newMdb.Machines = make([]mdb.Machine, 0, len(newMachines))
	for _, machine := range newMachines {
		newMdb.Machines = append(newMdb.Machines, machine)
	}
	return newMdb
}
<|endoftext|>


<commit_before>package srnd

import (
	"bufio"
	"errors"
	"io"
	"net"
)

var ErrSpamFilterNotEnabled = errors.New("spam filter access attempted when disabled")

type SpamFilter struct {
	addr    string
	enabled bool
}

func (sp *SpamFilter) Configure(c SpamConfig) {
	sp.enabled = c.enabled
	sp.addr = c.addr
}

func (sp *SpamFilter) Enabled() bool {
	return sp.enabled
}

func (sp *SpamFilter) Rewrite(msg io.Reader, out io.WriteCloser) error {
	var buff [65636]byte
	if !sp.Enabled() {
		return ErrSpamFilterNotEnabled
	}
	addr, err := net.ResolveTCPAddr("tcp", sp.addr)
	if err != nil {
		return err
	}
	c, err := net.DialTCP("tcp", nil, addr)
	if err != nil {
		return err
	}
	io.WriteString(c, "PROCESS SPAMC/1.5\r\nUser: nntpchan\r\n\r\n")
	io.CopyBuffer(c, msg, buff[:])
	c.CloseWrite()
	r := bufio.NewReader(c)
	r.ReadString(10)
	_, err = io.CopyBuffer(out, r, buff[:])
	c.Close()
	out.Close()
	return err
}
<commit_msg>fix sa hook<commit_after>package srnd

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"os/user"
)

var ErrSpamFilterNotEnabled = errors.New("spam filter access attempted when disabled")

type SpamFilter struct {
	addr    string
	enabled bool
}

func (sp *SpamFilter) Configure(c SpamConfig) {
	sp.enabled = c.enabled
	sp.addr = c.addr
}

func (sp *SpamFilter) Enabled() bool {
	return sp.enabled
}

func (sp *SpamFilter) Rewrite(msg io.Reader, out io.WriteCloser) error {
	var buff [65636]byte
	if !sp.Enabled() {
		return ErrSpamFilterNotEnabled
	}
	addr, err := net.ResolveTCPAddr("tcp", sp.addr)
	if err != nil {
		return err
	}
	c, err := net.DialTCP("tcp", nil, addr)
	if err != nil {
		return err
	}
	u, err := user.Current()
	if err != nil {
		return err
	}
	fmt.Fprintf(c, "PROCESS SPAMC/1.5\r\nUser: %s\r\n\r\n", u.Username)
	io.CopyBuffer(c, msg, buff[:])
	c.CloseWrite()
	r := bufio.NewReader(c)
	r.ReadString(10)
	_, err = io.CopyBuffer(out, r, buff[:])
	c.Close()
	out.Close()
	return err
}
<|endoftext|>


<commit_before>package lib

import (
	"bytes"
	"fmt"
	"strconv"

	"github.com/veino/logfan/parser"
	"github.com/veino/veino/config"
)

func parseConfigLocation(name string, options map[string]interface{}, pwd string, pickSections ...string) ([]config.Agent, error) {
	var locs Locations

	if _, ok := options["path"]; ok {
		locs.Add(options["path"].(string), pwd)
	} else if _, ok := options["url"]; ok {
		locs.Add(options["url"].(string), pwd)
	} else {
		return []config.Agent{}, fmt.Errorf("no location provided to get content from ; options=%v ", options)
	}

	return locs.Items[0].configAgentsWithOptions(options, pickSections...)
}

func buildAgents(content []byte, pwd string, pickSections ...string) ([]config.Agent, error) {
	var i int
	agentConfList := []config.Agent{}
	if len(pickSections) == 0 {
		pickSections = []string{"input", "filter", "output"}
	}

	p := parser.NewParser(bytes.NewReader(content))

	LSConfiguration, err := p.Parse()

	if err != nil {
		return agentConfList, err
	}

	outPorts := []config.Port{}

	if _, ok := LSConfiguration.Sections["input"]; ok && isInSlice("input", pickSections) {
		for pluginIndex := 0; pluginIndex < len(LSConfiguration.Sections["input"].Plugins); pluginIndex++ {
			plugin := LSConfiguration.Sections["input"].Plugins[pluginIndex]

			agents := buildInputAgents(plugin, pwd)

			agentConfList = append(agents, agentConfList...)
			outPort := config.Port{AgentID: agents[0].ID, PortNumber: 0}
			outPorts = append(outPorts, outPort)
		}
	}

	if _, ok := LSConfiguration.Sections["filter"]; ok && isInSlice("filter", pickSections) {
		if _, ok := LSConfiguration.Sections["filter"]; ok {
			for pluginIndex := 0; pluginIndex < len(LSConfiguration.Sections["filter"].Plugins); pluginIndex++ {
				var agents []config.Agent
				i++
				plugin := LSConfiguration.Sections["filter"].Plugins[pluginIndex]
				agents, outPorts = buildFilterAgents(plugin, outPorts, pwd)
				agentConfList = append(agents, agentConfList...)
			}
		}
	}

	if _, ok := LSConfiguration.Sections["output"]; ok && isInSlice("output", pickSections) {
		for pluginIndex := 0; pluginIndex < len(LSConfiguration.Sections["output"].Plugins); pluginIndex++ {
			var agents []config.Agent
			i++
			plugin := LSConfiguration.Sections["output"].Plugins[pluginIndex]
			agents = buildOutputAgents(plugin, outPorts, pwd)
			agentConfList = append(agents, agentConfList...)
		}
	}

	return agentConfList, nil
}

func buildInputAgents(plugin *parser.Plugin, pwd string) []config.Agent {

	var agent config.Agent
	agent = config.NewAgent()

	agent.Type = "input_" + plugin.Name
	agent.Label = fmt.Sprintf("%s", plugin.Name)
	agent.Buffer = 200
	agent.PoolSize = 1
	agent.Wd = pwd

	// Plugin configuration
	agent.Options = map[string]interface{}{}
	for _, setting := range plugin.Settings {
		agent.Options[setting.K] = setting.V
	}

	//todo : handle codec

	if plugin.Codec.Name != "" {
		agent.Options["codec"] = plugin.Codec.Name
	}

	// If agent is a "use"
	// build imported pipeline from path
	// connect import plugin Xsource to imported pipeline output
	if plugin.Name == "use" {
		fileConfigAgents, _ := parseConfigLocation("", agent.Options, pwd, "input", "filter")

		// add agent "use" - set use agent Source as last From FileConfigAgents
		inPort := config.Port{AgentID: fileConfigAgents[0].ID, PortNumber: 0}
		agent.XSources = append(agent.XSources, inPort)
		fileConfigAgents = append([]config.Agent{agent}, fileConfigAgents...)

		return fileConfigAgents
	}

	// interval can be a number, a string number or a cron string pattern
	interval := agent.Options["interval"]
	switch t := interval.(type) {
	case int, int8, int16, int32, int64:
		agent.Schedule = fmt.Sprintf("@every %ds", t)
	case string:
		if i, err := strconv.Atoi(t); err == nil {
			agent.Schedule = fmt.Sprintf("@every %ds", i)
		} else {
			agent.Schedule = t
		}
	}

	if workers, ok := agent.Options["workers"]; ok {
		switch t := workers.(type) {
		case int64:
			agent.PoolSize = int(t)
		case int32:
			agent.PoolSize = int(t)
		case string:
			if i, err := strconv.Atoi(t); err == nil {
				agent.PoolSize = i
			}
		}
	}
	// pp.Println("plugin-->", agent)
	return []config.Agent{agent}
}

func buildOutputAgents(plugin *parser.Plugin, lastOutPorts []config.Port, pwd string) []config.Agent {
	agent_list := []config.Agent{}

	var agent config.Agent
	agent = config.NewAgent()
	agent.Type = "output_" + plugin.Name
	agent.Label = fmt.Sprintf("%s", plugin.Name)
	agent.Buffer = 200
	agent.PoolSize = 1
	agent.Wd = pwd

	// Plugin configuration
	agent.Options = map[string]interface{}{}
	for _, setting := range plugin.Settings {
		agent.Options[setting.K] = setting.V
	}
	//todo : handle codec
	if plugin.Codec.Name != "" {
		agent.Options["codec"] = plugin.Codec.Name
	}
	// if its a use plugin
	// load filter and output parts of pipeline
	// connect pipeline Xsource to lastOutPorts
	// return pipelineagents with lastOutPorts intact
	// handle use plugin
	// If its a use agent
	// build the filter part of the pipeline
	// connect pipeline first agent Xsource to lastOutPorts output
	// return imported pipeline with its output
	if plugin.Name == "use" {
		fileConfigAgents, _ := parseConfigLocation("", agent.Options, pwd, "filter", "output")

		firstUsedAgent := &fileConfigAgents[len(fileConfigAgents)-1]
		for _, sourceport := range lastOutPorts {
			inPort := config.Port{AgentID: sourceport.AgentID, PortNumber: sourceport.PortNumber}
			firstUsedAgent.XSources = append(firstUsedAgent.XSources, inPort)
		}

		//specific to output
		return fileConfigAgents
	}

	// Plugin Sources
	agent.XSources = config.PortList{}
	for _, sourceport := range lastOutPorts {
		inPort := config.Port{AgentID: sourceport.AgentID, PortNumber: sourceport.PortNumber}
		agent.XSources = append(agent.XSources, inPort)
	}

	if plugin.Codec != nil {
		agent.Options["codec"] = plugin.Codec.Name
	}

	// Is this Plugin has conditional expressions ?
	if len(plugin.When) > 0 {
		// outPorts_when := []port{}
		// le plugin WHEn est $plugin
		agent.Options["expressions"] = map[int]string{}
		// Loop over expressions in correct order
		for expressionIndex := 0; expressionIndex < len(plugin.When); expressionIndex++ {
			when := plugin.When[expressionIndex]
			//	enregistrer l'expression dans la conf agent
			agent.Options["expressions"].(map[int]string)[expressionIndex] = when.Expression

			// recupérer le outport associé (expressionIndex)
			expressionOutPorts := []config.Port{
				{AgentID: agent.ID, PortNumber: expressionIndex},
			}

			// construire les plugins associés à l'expression
			// en utilisant le expressionOutPorts
			for pi := 0; pi < len(when.Plugins); pi++ {
				p := when.Plugins[pi]
				var agents []config.Agent

				// récupérer le dernier outport du plugin créé il devient expressionOutPorts
				agents = buildOutputAgents(p, expressionOutPorts, pwd)
				// ajoute l'agent à la liste des agents construits
				agent_list = append(agents, agent_list...)
			}
		}
	}

	// ajoute l'agent à la liste des agents construits
	agent_list = append([]config.Agent{agent}, agent_list...)
	return agent_list
}

func buildFilterAgents(plugin *parser.Plugin, lastOutPorts []config.Port, pwd string) ([]config.Agent, []config.Port) {

	agent_list := []config.Agent{}

	var agent config.Agent
	agent = config.NewAgent()

	agent.Type = plugin.Name
	agent.Label = fmt.Sprintf("%s", plugin.Name)
	agent.Buffer = 200
	agent.PoolSize = 2
	agent.Wd = pwd

	// Plugin configuration
	agent.Options = map[string]interface{}{}
	for _, setting := range plugin.Settings {
		agent.Options[setting.K] = setting.V
	}

	// handle use plugin
	// If its a use agent
	// build the filter part of the pipeline
	// connect pipeline first agent Xsource to lastOutPorts output
	// return imported pipeline with its output
	if plugin.Name == "use" {
		fileConfigAgents, _ := parseConfigLocation("", agent.Options, pwd, "filter")

		firstUsedAgent := &fileConfigAgents[len(fileConfigAgents)-1]
		for _, sourceport := range lastOutPorts {
			inPort := config.Port{AgentID: sourceport.AgentID, PortNumber: sourceport.PortNumber}
			firstUsedAgent.XSources = append(firstUsedAgent.XSources, inPort)
		}

		newOutPorts := []config.Port{
			{AgentID: fileConfigAgents[0].ID, PortNumber: 0},
		}
		return fileConfigAgents, newOutPorts
	}

	// interval can be a number, a string number or a cron string pattern
	interval := agent.Options["interval"]
	switch t := interval.(type) {
	case int, int8, int16, int32, int64:
		agent.Schedule = fmt.Sprintf("@every %ds", t)
	case string:
		if i, err := strconv.Atoi(t); err == nil {
			agent.Schedule = fmt.Sprintf("@every %ds", i)
		} else {
			agent.Schedule = t
		}
	}

	if workers, ok := agent.Options["workers"]; ok {
		switch t := workers.(type) {
		case int64:
			agent.PoolSize = int(t)
		case int32:
			agent.PoolSize = int(t)
		case string:
			if i, err := strconv.Atoi(t); err == nil {
				agent.PoolSize = i
			}
		}
	}

	// Plugin Sources
	agent.XSources = config.PortList{}
	for _, sourceport := range lastOutPorts {
		inPort := config.Port{AgentID: sourceport.AgentID, PortNumber: sourceport.PortNumber}
		agent.XSources = append(agent.XSources, inPort)
	}

	// By Default Agents output to port 0
	newOutPorts := []config.Port{
		{AgentID: agent.ID, PortNumber: 0},
	}

	// Is this Plugin has conditional expressions ?
	if len(plugin.When) > 0 {
		outPorts_when := []config.Port{}
		// le plugin WHEn est $plugin
		agent.Options["expressions"] = map[int]string{}
		elseOK := false
		// Loop over expressions in correct order
		for expressionIndex := 0; expressionIndex < len(plugin.When); expressionIndex++ {
			when := plugin.When[expressionIndex]
			//	enregistrer l'expression dans la conf agent
			agent.Options["expressions"].(map[int]string)[expressionIndex] = when.Expression
			if when.Expression == "true" {
				elseOK = true
			}
			// recupérer le outport associé (expressionIndex)
			expressionOutPorts := []config.Port{
				{AgentID: agent.ID, PortNumber: expressionIndex},
			}

			// construire les plugins associés à l'expression
			// en utilisant le outportA
			for pi := 0; pi < len(when.Plugins); pi++ {
				p := when.Plugins[pi]
				var agents []config.Agent
				// récupérer le dernier outport du plugin créé il devient outportA
				agents, expressionOutPorts = buildFilterAgents(p, expressionOutPorts, pwd)
				// ajoute l'agent à la liste des agents construits
				agent_list = append(agents, agent_list...)
			}
			// ajouter le dernier outportA de l'expression au outport final du when
			outPorts_when = append(expressionOutPorts, outPorts_when...)
		}
		newOutPorts = outPorts_when

		// If no else expression was found, insert one
		if elseOK == false {
			agent.Options["expressions"].(map[int]string)[len(agent.Options["expressions"].(map[int]string))] = "true"
			elseOutPorts := []config.Port{
				{AgentID: agent.ID, PortNumber: len(agent.Options["expressions"].(map[int]string)) - 1},
			}
			newOutPorts = append(elseOutPorts, newOutPorts...)
		}
	}

	// ajoute l'agent à la liste des agents construits
	agent_list = append([]config.Agent{agent}, agent_list...)
	return agent_list, newOutPorts
}

func isInSlice(needle string, candidates []string) bool {
	for _, symbolType := range candidates {
		if needle == symbolType {
			return true
		}
	}
	return false
}
<commit_msg>remove a comment<commit_after>package lib

import (
	"bytes"
	"fmt"
	"strconv"

	"github.com/veino/logfan/parser"
	"github.com/veino/veino/config"
)

func parseConfigLocation(name string, options map[string]interface{}, pwd string, pickSections ...string) ([]config.Agent, error) {
	var locs Locations

	if _, ok := options["path"]; ok {
		locs.Add(options["path"].(string), pwd)
	} else if _, ok := options["url"]; ok {
		locs.Add(options["url"].(string), pwd)
	} else {
		return []config.Agent{}, fmt.Errorf("no location provided to get content from ; options=%v ", options)
	}

	return locs.Items[0].configAgentsWithOptions(options, pickSections...)
}

func buildAgents(content []byte, pwd string, pickSections ...string) ([]config.Agent, error) {
	var i int
	agentConfList := []config.Agent{}
	if len(pickSections) == 0 {
		pickSections = []string{"input", "filter", "output"}
	}

	p := parser.NewParser(bytes.NewReader(content))

	LSConfiguration, err := p.Parse()

	if err != nil {
		return agentConfList, err
	}

	outPorts := []config.Port{}

	if _, ok := LSConfiguration.Sections["input"]; ok && isInSlice("input", pickSections) {
		for pluginIndex := 0; pluginIndex < len(LSConfiguration.Sections["input"].Plugins); pluginIndex++ {
			plugin := LSConfiguration.Sections["input"].Plugins[pluginIndex]

			agents := buildInputAgents(plugin, pwd)

			agentConfList = append(agents, agentConfList...)
			outPort := config.Port{AgentID: agents[0].ID, PortNumber: 0}
			outPorts = append(outPorts, outPort)
		}
	}

	if _, ok := LSConfiguration.Sections["filter"]; ok && isInSlice("filter", pickSections) {
		if _, ok := LSConfiguration.Sections["filter"]; ok {
			for pluginIndex := 0; pluginIndex < len(LSConfiguration.Sections["filter"].Plugins); pluginIndex++ {
				var agents []config.Agent
				i++
				plugin := LSConfiguration.Sections["filter"].Plugins[pluginIndex]
				agents, outPorts = buildFilterAgents(plugin, outPorts, pwd)
				agentConfList = append(agents, agentConfList...)
			}
		}
	}

	if _, ok := LSConfiguration.Sections["output"]; ok && isInSlice("output", pickSections) {
		for pluginIndex := 0; pluginIndex < len(LSConfiguration.Sections["output"].Plugins); pluginIndex++ {
			var agents []config.Agent
			i++
			plugin := LSConfiguration.Sections["output"].Plugins[pluginIndex]
			agents = buildOutputAgents(plugin, outPorts, pwd)
			agentConfList = append(agents, agentConfList...)
		}
	}

	return agentConfList, nil
}

func buildInputAgents(plugin *parser.Plugin, pwd string) []config.Agent {

	var agent config.Agent
	agent = config.NewAgent()

	agent.Type = "input_" + plugin.Name
	agent.Label = fmt.Sprintf("%s", plugin.Name)
	agent.Buffer = 200
	agent.PoolSize = 1
	agent.Wd = pwd

	// Plugin configuration
	agent.Options = map[string]interface{}{}
	for _, setting := range plugin.Settings {
		agent.Options[setting.K] = setting.V
	}

	//todo : handle codec

	if plugin.Codec.Name != "" {
		agent.Options["codec"] = plugin.Codec.Name
	}

	// If agent is a "use"
	// build imported pipeline from path
	// connect import plugin Xsource to imported pipeline output
	if plugin.Name == "use" {
		fileConfigAgents, _ := parseConfigLocation("", agent.Options, pwd, "input", "filter")

		// add agent "use" - set use agent Source as last From FileConfigAgents
		inPort := config.Port{AgentID: fileConfigAgents[0].ID, PortNumber: 0}
		agent.XSources = append(agent.XSources, inPort)
		fileConfigAgents = append([]config.Agent{agent}, fileConfigAgents...)

		return fileConfigAgents
	}

	// interval can be a number, a string number or a cron string pattern
	interval := agent.Options["interval"]
	switch t := interval.(type) {
	case int, int8, int16, int32, int64:
		agent.Schedule = fmt.Sprintf("@every %ds", t)
	case string:
		if i, err := strconv.Atoi(t); err == nil {
			agent.Schedule = fmt.Sprintf("@every %ds", i)
		} else {
			agent.Schedule = t
		}
	}

	if workers, ok := agent.Options["workers"]; ok {
		switch t := workers.(type) {
		case int64:
			agent.PoolSize = int(t)
		case int32:
			agent.PoolSize = int(t)
		case string:
			if i, err := strconv.Atoi(t); err == nil {
				agent.PoolSize = i
			}
		}
	}

	return []config.Agent{agent}
}

func buildOutputAgents(plugin *parser.Plugin, lastOutPorts []config.Port, pwd string) []config.Agent {
	agent_list := []config.Agent{}

	var agent config.Agent
	agent = config.NewAgent()
	agent.Type = "output_" + plugin.Name
	agent.Label = fmt.Sprintf("%s", plugin.Name)
	agent.Buffer = 200
	agent.PoolSize = 1
	agent.Wd = pwd

	// Plugin configuration
	agent.Options = map[string]interface{}{}
	for _, setting := range plugin.Settings {
		agent.Options[setting.K] = setting.V
	}
	//todo : handle codec
	if plugin.Codec.Name != "" {
		agent.Options["codec"] = plugin.Codec.Name
	}
	// if its a use plugin
	// load filter and output parts of pipeline
	// connect pipeline Xsource to lastOutPorts
	// return pipelineagents with lastOutPorts intact
	// handle use plugin
	// If its a use agent
	// build the filter part of the pipeline
	// connect pipeline first agent Xsource to lastOutPorts output
	// return imported pipeline with its output
	if plugin.Name == "use" {
		fileConfigAgents, _ := parseConfigLocation("", agent.Options, pwd, "filter", "output")

		firstUsedAgent := &fileConfigAgents[len(fileConfigAgents)-1]
		for _, sourceport := range lastOutPorts {
			inPort := config.Port{AgentID: sourceport.AgentID, PortNumber: sourceport.PortNumber}
			firstUsedAgent.XSources = append(firstUsedAgent.XSources, inPort)
		}

		//specific to output
		return fileConfigAgents
	}

	// Plugin Sources
	agent.XSources = config.PortList{}
	for _, sourceport := range lastOutPorts {
		inPort := config.Port{AgentID: sourceport.AgentID, PortNumber: sourceport.PortNumber}
		agent.XSources = append(agent.XSources, inPort)
	}

	if plugin.Codec != nil {
		agent.Options["codec"] = plugin.Codec.Name
	}

	// Is this Plugin has conditional expressions ?
	if len(plugin.When) > 0 {
		// outPorts_when := []port{}
		// le plugin WHEn est $plugin
		agent.Options["expressions"] = map[int]string{}
		// Loop over expressions in correct order
		for expressionIndex := 0; expressionIndex < len(plugin.When); expressionIndex++ {
			when := plugin.When[expressionIndex]
			//	enregistrer l'expression dans la conf agent
			agent.Options["expressions"].(map[int]string)[expressionIndex] = when.Expression

			// recupérer le outport associé (expressionIndex)
			expressionOutPorts := []config.Port{
				{AgentID: agent.ID, PortNumber: expressionIndex},
			}

			// construire les plugins associés à l'expression
			// en utilisant le expressionOutPorts
			for pi := 0; pi < len(when.Plugins); pi++ {
				p := when.Plugins[pi]
				var agents []config.Agent

				// récupérer le dernier outport du plugin créé il devient expressionOutPorts
				agents = buildOutputAgents(p, expressionOutPorts, pwd)
				// ajoute l'agent à la liste des agents construits
				agent_list = append(agents, agent_list...)
			}
		}
	}

	// ajoute l'agent à la liste des agents construits
	agent_list = append([]config.Agent{agent}, agent_list...)
	return agent_list
}

func buildFilterAgents(plugin *parser.Plugin, lastOutPorts []config.Port, pwd string) ([]config.Agent, []config.Port) {

	agent_list := []config.Agent{}

	var agent config.Agent
	agent = config.NewAgent()

	agent.Type = plugin.Name
	agent.Label = fmt.Sprintf("%s", plugin.Name)
	agent.Buffer = 200
	agent.PoolSize = 2
	agent.Wd = pwd

	// Plugin configuration
	agent.Options = map[string]interface{}{}
	for _, setting := range plugin.Settings {
		agent.Options[setting.K] = setting.V
	}

	// handle use plugin
	// If its a use agent
	// build the filter part of the pipeline
	// connect pipeline first agent Xsource to lastOutPorts output
	// return imported pipeline with its output
	if plugin.Name == "use" {
		fileConfigAgents, _ := parseConfigLocation("", agent.Options, pwd, "filter")

		firstUsedAgent := &fileConfigAgents[len(fileConfigAgents)-1]
		for _, sourceport := range lastOutPorts {
			inPort := config.Port{AgentID: sourceport.AgentID, PortNumber: sourceport.PortNumber}
			firstUsedAgent.XSources = append(firstUsedAgent.XSources, inPort)
		}

		newOutPorts := []config.Port{
			{AgentID: fileConfigAgents[0].ID, PortNumber: 0},
		}
		return fileConfigAgents, newOutPorts
	}

	// interval can be a number, a string number or a cron string pattern
	interval := agent.Options["interval"]
	switch t := interval.(type) {
	case int, int8, int16, int32, int64:
		agent.Schedule = fmt.Sprintf("@every %ds", t)
	case string:
		if i, err := strconv.Atoi(t); err == nil {
			agent.Schedule = fmt.Sprintf("@every %ds", i)
		} else {
			agent.Schedule = t
		}
	}

	if workers, ok := agent.Options["workers"]; ok {
		switch t := workers.(type) {
		case int64:
			agent.PoolSize = int(t)
		case int32:
			agent.PoolSize = int(t)
		case string:
			if i, err := strconv.Atoi(t); err == nil {
				agent.PoolSize = i
			}
		}
	}

	// Plugin Sources
	agent.XSources = config.PortList{}
	for _, sourceport := range lastOutPorts {
		inPort := config.Port{AgentID: sourceport.AgentID, PortNumber: sourceport.PortNumber}
		agent.XSources = append(agent.XSources, inPort)
	}

	// By Default Agents output to port 0
	newOutPorts := []config.Port{
		{AgentID: agent.ID, PortNumber: 0},
	}

	// Is this Plugin has conditional expressions ?
	if len(plugin.When) > 0 {
		outPorts_when := []config.Port{}
		// le plugin WHEn est $plugin
		agent.Options["expressions"] = map[int]string{}
		elseOK := false
		// Loop over expressions in correct order
		for expressionIndex := 0; expressionIndex < len(plugin.When); expressionIndex++ {
			when := plugin.When[expressionIndex]
			//	enregistrer l'expression dans la conf agent
			agent.Options["expressions"].(map[int]string)[expressionIndex] = when.Expression
			if when.Expression == "true" {
				elseOK = true
			}
			// recupérer le outport associé (expressionIndex)
			expressionOutPorts := []config.Port{
				{AgentID: agent.ID, PortNumber: expressionIndex},
			}

			// construire les plugins associés à l'expression
			// en utilisant le outportA
			for pi := 0; pi < len(when.Plugins); pi++ {
				p := when.Plugins[pi]
				var agents []config.Agent
				// récupérer le dernier outport du plugin créé il devient outportA
				agents, expressionOutPorts = buildFilterAgents(p, expressionOutPorts, pwd)
				// ajoute l'agent à la liste des agents construits
				agent_list = append(agents, agent_list...)
			}
			// ajouter le dernier outportA de l'expression au outport final du when
			outPorts_when = append(expressionOutPorts, outPorts_when...)
		}
		newOutPorts = outPorts_when

		// If no else expression was found, insert one
		if elseOK == false {
			agent.Options["expressions"].(map[int]string)[len(agent.Options["expressions"].(map[int]string))] = "true"
			elseOutPorts := []config.Port{
				{AgentID: agent.ID, PortNumber: len(agent.Options["expressions"].(map[int]string)) - 1},
			}
			newOutPorts = append(elseOutPorts, newOutPorts...)
		}
	}

	// ajoute l'agent à la liste des agents construits
	agent_list = append([]config.Agent{agent}, agent_list...)
	return agent_list, newOutPorts
}

func isInSlice(needle string, candidates []string) bool {
	for _, symbolType := range candidates {
		if needle == symbolType {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package kit

import (
	"fmt"
	"go/format"

	"github.com/cihangir/gene/generators/common"
	"github.com/cihangir/schema"
)

// GenerateKitWorker generates the worker system for base kit
func GenerateKitWorker(context *common.Context, s *schema.Schema) ([]common.Output, error) {
	var outputs []common.Output

	for name, template := range templates {
		path := fmt.Sprintf(
			"%s/kitworker/%s.go",
			context.Config.Target,
			name,
		)

		api, err := format.Source([]byte(template))
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content: api,
			Path:    path,
		})
	}

	return outputs, nil

}

var templates = map[string]string{
	"instrumenting": `package kitworker

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

	"github.com/go-kit/kit/endpoint"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/metrics"
	"golang.org/x/net/context"
)

// DefaultMiddlewares provides bare bones for default middlewares with
// requestLatency, requestCount and requestLogging
func DefaultMiddlewares(method string, requestCount metrics.Counter, requestLatency metrics.Histogram, logger log.Logger) endpoint.Middleware {
	return endpoint.Chain(
		RequestLatencyMiddleware(method, requestLatency),
		RequestCountMiddleware(method, requestCount),
		RequestLoggingMiddleware(method, logger),
	)
}

// RequestCountMiddleware prepares a request counter endpoint.Middleware for
// package wide usage
func RequestCountMiddleware(method string, requestCount metrics.Counter) endpoint.Middleware {
	return func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, request interface{}) (response interface{}, err error) {
			defer func() {
				requestCount.With("method", method).With("error", fmt.Sprintf("%v", err)).Add(1)
			}()

			response, err = next(ctx, request)
			return
		}
	}
}

// RequestLatencyMiddleware prepares a request latency calculator
// endpoint.Middleware for package wide usage
func RequestLatencyMiddleware(method string, requestLatency metrics.Histogram) endpoint.Middleware {
	return func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, request interface{}) (response interface{}, err error) {
			defer func(begin time.Time) {
				requestLatency.With("method", method).With("error", fmt.Sprintf("%v", err)).Observe(float64(time.Since(begin)))
			}(time.Now())
			response, err = next(ctx, request)
			return
		}
	}
}

// RequestLoggingMiddleware prepares a request logger endpoint.Middleware for
// package wide usage
func RequestLoggingMiddleware(method string, logger log.Logger) endpoint.Middleware {
	return func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, request interface{}) (response interface{}, err error) {
			defer func(begin time.Time) {
				input, _ := json.Marshal(request)
				output, _ := json.Marshal(response)
				_ = logger.Log(
					"method", method,
					"input", string(input),
					"output", string(output),
					"err", err,
					"took", time.Since(begin),
				)
			}(time.Now())
			response, err = next(ctx, request)
			return
		}
	}
}
`,

	"client": `package kitworker

import (
	"net/url"
	"strings"

	"github.com/go-kit/kit/circuitbreaker"
	"github.com/go-kit/kit/endpoint"
	kitratelimit "github.com/go-kit/kit/ratelimit"
	"github.com/go-kit/kit/sd"
	"github.com/go-kit/kit/sd/lb"
	"github.com/go-kit/kit/tracing/opentracing"
	httptransport "github.com/go-kit/kit/transport/http"
	jujuratelimit "github.com/juju/ratelimit"
	stdopentracing "github.com/opentracing/opentracing-go"
	"github.com/sony/gobreaker"
)

// LoadBalancerF
type LoadBalancerF func(factory sd.Factory) lb.Balancer

// ClientOption holds the required parameters for configuring a client
type ClientOption struct {
	// Host holds the host's name
	Host string

	// ZipkinCollector holds the collector for zipkin tracing
	Tracer stdopentracing.Tracer

	// DisableCircuitBreaker disables circuit breaking functionality
	DisableCircuitBreaker bool

	// CircuitBreaker holds the custom circuit breaker, if not set a default one
	// will be created with default settings
	CircuitBreaker *gobreaker.CircuitBreaker

	// DisableRateLimiter disables rate limiting functionality
	DisableRateLimiter bool

	// QPS holds the configration parameter for rate limiting outgoing requests
	// to remote client. Must be set othervise all requests will be blocked
	// unless rate limiting is disabled
	QPS int

	// RateLimiter holds the custom rate limiter, if not set a default one will be created automatically
	RateLimiter *jujuratelimit.Bucket

	// TransportOpts holds custom httptransport.ClientOption array will be
	// appended to the end of the autogenerated ClientOptions
	TransportOpts []httptransport.ClientOption

	// Middlewares holds custom endpoint.Middleware array will be appended to
	// the end of the autogenerated Middlewares
	Middlewares []endpoint.Middleware

	// LoadBalancerCreator creates the loadbalancing strategy after getting the factory
	LoadBalancerCreator LoadBalancerF
}

// Configure prepares middlewares and clientOptions from the client options
//
// If required:
//   Adds circuitbreaker from "github.com/sony/gobreaker"
//   Adds ratelimiting from  "github.com/juju/ratelimit"
//   Adds request tracing from "github.com/go-kit/kit/tracing/zipkin"
func (c ClientOption) Configure(moduleName, funcName string) ([]endpoint.Middleware, []httptransport.ClientOption) {
	var transportOpts []httptransport.ClientOption
	var middlewares []endpoint.Middleware

	// if circuit braker is not disabled, add it as a middleware
	if !c.DisableCircuitBreaker {
		cb := c.CircuitBreaker

		if c.CircuitBreaker == nil {
			// create a default circuit breaker
			cb = gobreaker.NewCircuitBreaker(gobreaker.Settings{})
		}

		middlewares = append(middlewares, circuitbreaker.Gobreaker(cb))
	}

	// if rate limiter is not disabled, add it as a middleware
	if !c.DisableRateLimiter {
		rateLimiter := c.RateLimiter

		if c.RateLimiter == nil {
			// create a default rate limiter
			rateLimiter = jujuratelimit.NewBucketWithRate(float64(c.QPS), int64(c.QPS))
		}

		middlewares = append(middlewares, kitratelimit.NewTokenBucketLimiter(rateLimiter))
	}

	// enable tracing if required
	if c.Tracer != nil {
		middlewares = append(middlewares, opentracing.TraceServer(c.Tracer, funcName))
	}

	// If any custom middlewares are passed include them
	if len(c.Middlewares) > 0 {
		middlewares = append(middlewares, c.Middlewares...)
	}

	// If any client options are passed include them in client creation
	if len(c.TransportOpts) > 0 {
		transportOpts = append(transportOpts, c.TransportOpts...)
	}

	return middlewares, transportOpts
}

// CreateProxyURL creates an URL as proxy URL
func CreateProxyURL(instance, endpoint string) *url.URL {
	if !strings.HasPrefix(instance, "http") {
		instance = "http://" + instance
	}
	u, err := url.Parse(instance)
	if err != nil {
		panic(err)
	}
	if u.Path == "" {
		u.Path = endpoint
	}

	return u
}
`,

	"server": `package kitworker

import (
	"github.com/go-kit/kit/endpoint"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/metrics"
	"github.com/go-kit/kit/tracing/opentracing"
	httptransport "github.com/go-kit/kit/transport/http"
	stdopentracing "github.com/opentracing/opentracing-go"
)

// ServerOption holds the required parameters for configuring a server
type ServerOption struct {
	// Host holds the host's name
	Host string

	// Tracer holds the collector for zipkin tracing
	Tracer stdopentracing.Tracer

	// LogErrors configures whether server should log error responses or not
	LogErrors bool

	// LogRequests configures if the server should log incoming requests or not
	LogRequests bool

	// Latency holds the metric metric for request latency metric collection, if
	// not set Latency metrics will not be collected
	Latency metrics.Histogram

	// Counter holds the metrics.Counter metric for request count metric
	// collection, if not set RequestCountMetrics will not be collected
	Counter metrics.Counter

	// ServerOptions holds custom httptransport.ServerOption array, will be
	// appended to the end of the autogenerated
	ServerOptions []httptransport.ServerOption

	// Middlewares holds custom endpoint.Middleware array will be appended to
	// the end of the autogenerated Middlewares
	Middlewares []endpoint.Middleware
}

// Configure prepares middlewares and serverOptions from the client options
//
// If required:
//   Adds RequestLatencyMiddleware
//   Adds RequestCountMiddleware
//   Adds RequestLoggingMiddleware
//   Adds Zipkin Tracing
//   Adds httptransport.ServerErrorLogger
func (s ServerOption) Configure(moduleName, funcName string, logger log.Logger) ([]endpoint.Middleware, []httptransport.ServerOption) {

	var serverOpts []httptransport.ServerOption
	var middlewares []endpoint.Middleware

	if s.Latency != nil {
		middlewares = append(middlewares, RequestLatencyMiddleware(funcName, s.Latency))
	}

	if s.Counter != nil {
		middlewares = append(middlewares, RequestCountMiddleware(funcName, s.Counter))
	}

	if s.LogRequests {
		middlewares = append(middlewares, RequestLoggingMiddleware(funcName, logger))
	}

	// enable tracing if required
	if s.Tracer != nil {
		middlewares = append(middlewares, opentracing.TraceServer(s.Tracer, funcName))
	}

	// log server errors
	if s.LogErrors {
		serverOpts = append(serverOpts, httptransport.ServerErrorLogger(logger))
	}

	// If any custom middlewares are passed include them
	if len(s.Middlewares) > 0 {
		middlewares = append(middlewares, s.Middlewares...)
	}

	// If any server options are passed include them in server creation
	if len(s.ServerOptions) > 0 {
		serverOpts = append(serverOpts, s.ServerOptions...)
	}

	return middlewares, serverOpts
}
`,
}
<commit_msg>generator: fix docs and add http tracing<commit_after>package kit

import (
	"fmt"
	"go/format"

	"github.com/cihangir/gene/generators/common"
	"github.com/cihangir/schema"
)

// GenerateKitWorker generates the worker system for base kit
func GenerateKitWorker(context *common.Context, s *schema.Schema) ([]common.Output, error) {
	var outputs []common.Output

	for name, template := range templates {
		path := fmt.Sprintf(
			"%s/kitworker/%s.go",
			context.Config.Target,
			name,
		)

		api, err := format.Source([]byte(template))
		if err != nil {
			return nil, err
		}

		outputs = append(outputs, common.Output{
			Content: api,
			Path:    path,
		})
	}

	return outputs, nil

}

var templates = map[string]string{
	"instrumenting": `package kitworker

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

	"github.com/go-kit/kit/endpoint"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/metrics"
	"golang.org/x/net/context"
)

// DefaultMiddlewares provides bare bones for default middlewares with
// requestLatency, requestCount and requestLogging
func DefaultMiddlewares(method string, requestCount metrics.Counter, requestLatency metrics.Histogram, logger log.Logger) endpoint.Middleware {
	return endpoint.Chain(
		RequestLatencyMiddleware(method, requestLatency),
		RequestCountMiddleware(method, requestCount),
		RequestLoggingMiddleware(method, logger),
	)
}

// RequestCountMiddleware prepares a request counter endpoint.Middleware for
// package wide usage
func RequestCountMiddleware(method string, requestCount metrics.Counter) endpoint.Middleware {
	return func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, request interface{}) (response interface{}, err error) {
			defer func() {
				requestCount.With("method", method).With("error", fmt.Sprintf("%v", err)).Add(1)
			}()

			response, err = next(ctx, request)
			return
		}
	}
}

// RequestLatencyMiddleware prepares a request latency calculator
// endpoint.Middleware for package wide usage
func RequestLatencyMiddleware(method string, requestLatency metrics.Histogram) endpoint.Middleware {
	return func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, request interface{}) (response interface{}, err error) {
			defer func(begin time.Time) {
				requestLatency.With("method", method).With("error", fmt.Sprintf("%v", err)).Observe(float64(time.Since(begin)))
			}(time.Now())
			response, err = next(ctx, request)
			return
		}
	}
}

// RequestLoggingMiddleware prepares a request logger endpoint.Middleware for
// package wide usage
func RequestLoggingMiddleware(method string, logger log.Logger) endpoint.Middleware {
	return func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, request interface{}) (response interface{}, err error) {
			defer func(begin time.Time) {
				input, _ := json.Marshal(request)
				output, _ := json.Marshal(response)
				_ = logger.Log(
					"method", method,
					"input", string(input),
					"output", string(output),
					"err", err,
					"took", time.Since(begin),
				)
			}(time.Now())
			response, err = next(ctx, request)
			return
		}
	}
}
`,

	"client": `package kitworker

import (
	"net/url"
	"strings"

	"github.com/go-kit/kit/circuitbreaker"
	"github.com/go-kit/kit/endpoint"
	kitratelimit "github.com/go-kit/kit/ratelimit"
	"github.com/go-kit/kit/sd"
	"github.com/go-kit/kit/sd/lb"
	"github.com/go-kit/kit/tracing/opentracing"
	httptransport "github.com/go-kit/kit/transport/http"
	jujuratelimit "github.com/juju/ratelimit"
	stdopentracing "github.com/opentracing/opentracing-go"
	"github.com/sony/gobreaker"
)

// LoadBalancerF
type LoadBalancerF func(factory sd.Factory) lb.Balancer

// ClientOption holds the required parameters for configuring a client
type ClientOption struct {
	// Host holds the host's name
	Host string

	// Tracer holds the collector for tracing
	Tracer stdopentracing.Tracer

	// DisableCircuitBreaker disables circuit breaking functionality
	DisableCircuitBreaker bool

	// CircuitBreaker holds the custom circuit breaker, if not set a default one
	// will be created with default settings
	CircuitBreaker *gobreaker.CircuitBreaker

	// DisableRateLimiter disables rate limiting functionality
	DisableRateLimiter bool

	// QPS holds the configration parameter for rate limiting outgoing requests
	// to remote client. Must be set othervise all requests will be blocked
	// unless rate limiting is disabled
	QPS int

	// RateLimiter holds the custom rate limiter, if not set a default one will be created automatically
	RateLimiter *jujuratelimit.Bucket

	// TransportOpts holds custom httptransport.ClientOption array will be
	// appended to the end of the autogenerated ClientOptions
	TransportOpts []httptransport.ClientOption

	// Middlewares holds custom endpoint.Middleware array will be appended to
	// the end of the autogenerated Middlewares
	Middlewares []endpoint.Middleware

	// LoadBalancerCreator creates the loadbalancing strategy after getting the factory
	LoadBalancerCreator LoadBalancerF
}

// Configure prepares middlewares and clientOptions from the client options
//
// If required:
//   Adds circuitbreaker from "github.com/sony/gobreaker"
//   Adds ratelimiting from  "github.com/juju/ratelimit"
func (c ClientOption) Configure(moduleName, funcName string) ([]endpoint.Middleware, []httptransport.ClientOption) {
	var transportOpts []httptransport.ClientOption
	var middlewares []endpoint.Middleware

	// if circuit braker is not disabled, add it as a middleware
	if !c.DisableCircuitBreaker {
		cb := c.CircuitBreaker

		if c.CircuitBreaker == nil {
			// create a default circuit breaker
			cb = gobreaker.NewCircuitBreaker(gobreaker.Settings{})
		}

		middlewares = append(middlewares, circuitbreaker.Gobreaker(cb))
	}

	// if rate limiter is not disabled, add it as a middleware
	if !c.DisableRateLimiter {
		rateLimiter := c.RateLimiter

		if c.RateLimiter == nil {
			// create a default rate limiter
			rateLimiter = jujuratelimit.NewBucketWithRate(float64(c.QPS), int64(c.QPS))
		}

		middlewares = append(middlewares, kitratelimit.NewTokenBucketLimiter(rateLimiter))
	}

	// enable tracing if required
	if s.Tracer != nil {
		middlewares = append(middlewares, opentracing.TraceServer(s.Tracer, funcName))
		serverOpts = append(serverOpts, httptransport.ServerBefore(
			opentracing.FromHTTPRequest(s.Tracer, funcName, logger),
		))
	}

	// If any custom middlewares are passed include them
	if len(c.Middlewares) > 0 {
		middlewares = append(middlewares, c.Middlewares...)
	}

	// If any client options are passed include them in client creation
	if len(c.TransportOpts) > 0 {
		transportOpts = append(transportOpts, c.TransportOpts...)
	}

	return middlewares, transportOpts
}

// CreateProxyURL creates an URL as proxy URL
func CreateProxyURL(instance, endpoint string) *url.URL {
	if !strings.HasPrefix(instance, "http") {
		instance = "http://" + instance
	}
	u, err := url.Parse(instance)
	if err != nil {
		panic(err)
	}
	if u.Path == "" {
		u.Path = endpoint
	}

	return u
}
`,

	"server": `package kitworker

import (
	"github.com/go-kit/kit/endpoint"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/metrics"
	"github.com/go-kit/kit/tracing/opentracing"
	httptransport "github.com/go-kit/kit/transport/http"
	stdopentracing "github.com/opentracing/opentracing-go"
)

// ServerOption holds the required parameters for configuring a server
type ServerOption struct {
	// Host holds the host's name
	Host string

	// Tracer holds the collector for tracing
	Tracer stdopentracing.Tracer

	// LogErrors configures whether server should log error responses or not
	LogErrors bool

	// LogRequests configures if the server should log incoming requests or not
	LogRequests bool

	// Latency holds the metric metric for request latency metric collection, if
	// not set Latency metrics will not be collected
	Latency metrics.Histogram

	// Counter holds the metrics.Counter metric for request count metric
	// collection, if not set RequestCountMetrics will not be collected
	Counter metrics.Counter

	// ServerOptions holds custom httptransport.ServerOption array, will be
	// appended to the end of the autogenerated
	ServerOptions []httptransport.ServerOption

	// Middlewares holds custom endpoint.Middleware array will be appended to
	// the end of the autogenerated Middlewares
	Middlewares []endpoint.Middleware
}

// Configure prepares middlewares and serverOptions from the client options
//
// If required:
//   Adds RequestLatencyMiddleware
//   Adds RequestCountMiddleware
//   Adds RequestLoggingMiddleware
//   Adds httptransport.ServerErrorLogger
func (s ServerOption) Configure(moduleName, funcName string, logger log.Logger) ([]endpoint.Middleware, []httptransport.ServerOption) {

	var serverOpts []httptransport.ServerOption
	var middlewares []endpoint.Middleware

	if s.Latency != nil {
		middlewares = append(middlewares, RequestLatencyMiddleware(funcName, s.Latency))
	}

	if s.Counter != nil {
		middlewares = append(middlewares, RequestCountMiddleware(funcName, s.Counter))
	}

	if s.LogRequests {
		middlewares = append(middlewares, RequestLoggingMiddleware(funcName, logger))
	}

	// enable tracing if required
	if s.Tracer != nil {
		middlewares = append(middlewares, opentracing.TraceServer(s.Tracer, funcName))
	}

	// log server errors
	if s.LogErrors {
		serverOpts = append(serverOpts, httptransport.ServerErrorLogger(logger))
	}

	// If any custom middlewares are passed include them
	if len(s.Middlewares) > 0 {
		middlewares = append(middlewares, s.Middlewares...)
	}

	// If any server options are passed include them in server creation
	if len(s.ServerOptions) > 0 {
		serverOpts = append(serverOpts, s.ServerOptions...)
	}

	return middlewares, serverOpts
}
`,
}
<|endoftext|>


<commit_before>// Copyright 2020 Anapaya Systems
//
// 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 certs

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/x509"
	"fmt"
	"sort"
	"strings"

	"github.com/spf13/cobra"

	"github.com/scionproto/scion/go/lib/scrypto/cppki"
	"github.com/scionproto/scion/go/lib/serrors"
	"github.com/scionproto/scion/go/pkg/command"
)

var certTypes = map[string]cppki.CertType{
	cppki.Root.String():      cppki.Root,
	cppki.CA.String():        cppki.CA,
	cppki.AS.String():        cppki.AS,
	cppki.Sensitive.String(): cppki.Sensitive,
	cppki.Regular.String():   cppki.Regular,
}

func getTypes() []string {
	options := make([]string, 0, len(certTypes)+1)
	for k := range certTypes {
		options = append(options, k)
	}
	options = append(options, "any")
	options = append(options, "chain")
	sort.Strings(options)
	return options
}

func Cmd(pather command.Pather) *cobra.Command {
	cmd := &cobra.Command{
		Use:     "certificate",
		Aliases: []string{"cert", "certs"},
		Short:   "Manage certificates for the SCION control plane PKI.",
	}
	joined := command.Join(pather, cmd)
	cmd.AddCommand(
		newCreateCmd(joined),
		newValidateCmd(joined),
		newVerifyCmd(joined),
		newRenewCmd(joined),
	)
	return cmd
}

func newValidateCmd(pather command.Pather) *cobra.Command {
	var flags struct {
		certType string
	}

	cmd := &cobra.Command{
		Use:   "validate",
		Short: "Validate a SCION cert according to its type",
		Long: `'validate' checks if the certificate is valid and of the specified type.

In case the 'any' type is specified, this command attempts to identify what type
a certificate is and validates it accordingly. The identified type is stated in
the output.
`,
		Example: fmt.Sprintf(`  %[1]s validate --type cp-root /tmp/certs/cp-root.crt
  %[1]s validate --type any /tmp/certs/cp-root.crt`, pather.CommandPath()),
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			expectedType, checkType := certTypes[flags.certType]
			if !checkType && (flags.certType != "any" && flags.certType != "chain") {
				return serrors.New("invalid type flag", "type", flags.certType)
			}
			cmd.SilenceUsage = true

			filename := args[0]
			certs, err := cppki.ReadPEMCerts(filename)
			if err != nil {
				return err
			}
			if flags.certType == "chain" || len(certs) != 1 && flags.certType == "any" {
				if err := validateChain(certs); err != nil {
					return err
				}
				fmt.Printf("Valid certificate chain: %q\n", filename)

			} else {
				ct, err := validateCert(certs, expectedType, checkType)
				if err != nil {
					return err
				}
				fmt.Printf("Valid %s certificate: %q\n", ct, filename)
			}
			return nil
		},
	}

	cmd.Flags().StringVar(&flags.certType, "type", "",
		fmt.Sprintf("type of cert (%s) (required)", strings.Join(getTypes(), "|")))
	cmd.MarkFlagRequired("type")

	return cmd
}

func validateChain(certs []*x509.Certificate) error {
	if err := cppki.ValidateChain(certs); err != nil {
		return err
	}
	checkAlgorithm(certs[0])
	checkAlgorithm(certs[1])
	return nil
}

func validateCert(
	certs []*x509.Certificate,
	expectedType cppki.CertType,
	checkType bool,
) (cppki.CertType, error) {

	if len(certs) > 1 {
		return cppki.Invalid, serrors.New("file with multiple certificates not supported")
	}
	cert := certs[0]
	ct, err := cppki.ValidateCert(cert)
	if err != nil {
		return cppki.Invalid, err
	}
	if checkType && expectedType != ct {
		return cppki.Invalid, serrors.New("wrong certificate type",
			"expected", expectedType,
			"actual", ct,
		)
	}
	if ct == cppki.Root || ct == cppki.Regular || ct == cppki.Sensitive {
		checkAlgorithm(cert)
	}
	return ct, nil
}

func checkAlgorithm(cert *x509.Certificate) {
	if cert.PublicKeyAlgorithm != x509.ECDSA {
		return
	}

	pub, ok := cert.PublicKey.(*ecdsa.PublicKey)
	if !ok {
		return
	}
	expected := map[elliptic.Curve]x509.SignatureAlgorithm{
		elliptic.P256(): x509.ECDSAWithSHA256,
		elliptic.P384(): x509.ECDSAWithSHA384,
		elliptic.P521(): x509.ECDSAWithSHA512,
	}[pub.Curve]
	if expected != cert.SignatureAlgorithm {
		fmt.Printf("WARNING: Signature with %s curve should use %s instead of %s\n",
			pub.Curve.Params().Name, expected, cert.SignatureAlgorithm)
	}
}
<commit_msg>scion-pki: add --check-time to certificate validate<commit_after>// Copyright 2020 Anapaya Systems
//
// 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 certs

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/x509"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/spf13/cobra"

	"github.com/scionproto/scion/go/lib/scrypto/cppki"
	"github.com/scionproto/scion/go/lib/serrors"
	"github.com/scionproto/scion/go/pkg/app/flag"
	"github.com/scionproto/scion/go/pkg/command"
)

var certTypes = map[string]cppki.CertType{
	cppki.Root.String():      cppki.Root,
	cppki.CA.String():        cppki.CA,
	cppki.AS.String():        cppki.AS,
	cppki.Sensitive.String(): cppki.Sensitive,
	cppki.Regular.String():   cppki.Regular,
}

func getTypes() []string {
	options := make([]string, 0, len(certTypes)+1)
	for k := range certTypes {
		options = append(options, k)
	}
	options = append(options, "any")
	options = append(options, "chain")
	sort.Strings(options)
	return options
}

func Cmd(pather command.Pather) *cobra.Command {
	cmd := &cobra.Command{
		Use:     "certificate",
		Aliases: []string{"cert", "certs"},
		Short:   "Manage certificates for the SCION control plane PKI.",
	}
	joined := command.Join(pather, cmd)
	cmd.AddCommand(
		newCreateCmd(joined),
		newValidateCmd(joined),
		newVerifyCmd(joined),
		newRenewCmd(joined),
	)
	return cmd
}

func newValidateCmd(pather command.Pather) *cobra.Command {
	now := time.Now()
	var flags struct {
		certType    string
		checkTime   bool
		currentTime flag.Time
	}
	flags.currentTime = flag.Time{
		Time:    now,
		Current: now,
	}
	cmd := &cobra.Command{
		Use:   "validate",
		Short: "Validate a SCION cert according to its type",
		Long: `'validate' checks if the certificate is valid and of the specified type.

In case the 'any' type is specified, this command attempts to identify what type
a certificate is and validates it accordingly. The identified type is stated in
the output.

By default, the command does not check that the certificate is in its validity
period. This can be enabled by specifying the --check-time flag.
`,
		Example: fmt.Sprintf(`  %[1]s validate --type cp-root /tmp/certs/cp-root.crt
  %[1]s validate --type any /tmp/certs/cp-root.crt`, pather.CommandPath()),
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			expectedType, checkType := certTypes[flags.certType]
			if !checkType && (flags.certType != "any" && flags.certType != "chain") {
				return serrors.New("invalid type flag", "type", flags.certType)
			}
			cmd.SilenceUsage = true

			filename := args[0]
			certs, err := cppki.ReadPEMCerts(filename)
			if err != nil {
				return err
			}
			if flags.checkTime {
				validity := cppki.Validity{
					NotBefore: certs[0].NotBefore,
					NotAfter:  certs[0].NotAfter,
				}
				if current := flags.currentTime.Time; !validity.Contains(current) {
					return serrors.New("time not covered by certificate",
						"current_time", current,
						"validity.not_before", validity.NotBefore,
						"validity.not_after", validity.NotAfter,
					)
				}
			}
			if flags.certType == "chain" || len(certs) != 1 && flags.certType == "any" {
				if err := validateChain(certs); err != nil {
					return err
				}
				fmt.Printf("Valid certificate chain: %q\n", filename)

			} else {
				ct, err := validateCert(certs, expectedType, checkType)
				if err != nil {
					return err
				}
				fmt.Printf("Valid %s certificate: %q\n", ct, filename)
			}
			return nil
		},
	}

	cmd.Flags().StringVar(&flags.certType, "type", "",
		fmt.Sprintf("type of cert (%s) (required)", strings.Join(getTypes(), "|")),
	)
	cmd.Flags().BoolVar(&flags.checkTime, "check-time", false,
		"Check that the certificate covers the current time.",
	)
	cmd.Flags().Var(&flags.currentTime, "current-time",
		`The time that needs to be covered by the certificate.
Can either be a timestamp or an offset.

If the value is a timestamp, it is expected to either be an RFC 3339 formatted
timestamp or a unix timestamp. If the value is a duration, it is used as the
offset from the current time.`,
	)
	cmd.MarkFlagRequired("type")

	return cmd
}

func validateChain(certs []*x509.Certificate) error {
	if err := cppki.ValidateChain(certs); err != nil {
		return err
	}
	checkAlgorithm(certs[0])
	checkAlgorithm(certs[1])
	return nil
}

func validateCert(
	certs []*x509.Certificate,
	expectedType cppki.CertType,
	checkType bool,
) (cppki.CertType, error) {

	if len(certs) > 1 {
		return cppki.Invalid, serrors.New("file with multiple certificates not supported")
	}
	cert := certs[0]
	ct, err := cppki.ValidateCert(cert)
	if err != nil {
		return cppki.Invalid, err
	}
	if checkType && expectedType != ct {
		return cppki.Invalid, serrors.New("wrong certificate type",
			"expected", expectedType,
			"actual", ct,
		)
	}
	if ct == cppki.Root || ct == cppki.Regular || ct == cppki.Sensitive {
		checkAlgorithm(cert)
	}
	return ct, nil
}

func checkAlgorithm(cert *x509.Certificate) {
	if cert.PublicKeyAlgorithm != x509.ECDSA {
		return
	}

	pub, ok := cert.PublicKey.(*ecdsa.PublicKey)
	if !ok {
		return
	}
	expected := map[elliptic.Curve]x509.SignatureAlgorithm{
		elliptic.P256(): x509.ECDSAWithSHA256,
		elliptic.P384(): x509.ECDSAWithSHA384,
		elliptic.P521(): x509.ECDSAWithSHA512,
	}[pub.Curve]
	if expected != cert.SignatureAlgorithm {
		fmt.Printf("WARNING: Signature with %s curve should use %s instead of %s\n",
			pub.Curve.Params().Name, expected, cert.SignatureAlgorithm)
	}
}
<|endoftext|>