<commit_before>package msg

/* Constants indicating the HTTP Content-type for each proto/json type in this package */
const (
	StatusContentType           string = "application/vnd.zvelo.status"
	StatusContentTypeJSON       string = "application/vnd.zvelo.status+json"
	QueryContentType            string = "application/vnd.zvelo.query"
	QueryContentTypeJSON        string = "application/vnd.zvelo.query+json"
	ContentQueryContentType     string = "application/vnd.zvelo.content-query"
	ContentQueryContentTypeJSON string = "application/vnd.zvelo.content-query+json"
	QueryReplyContentType       string = "application/vnd.zvelo.query-reply"
	QueryReplyContentTypeJSON   string = "application/vnd.zvelo.query-reply+json"
	QueryResultContentType      string = "application/vnd.zvelo.query-result"
	QueryResultContentTypeJSON  string = "application/vnd.zvelo.query-result+json"

	StreamResultsContentType     string = "application/vnd.zvelo.stream-results"
	StreamResultsContentTypeJSON string = "application/vnd.zvelo.stream-results+json"
	StreamReplyContentType       string = "application/vnd.zvelo.stream-reply"
	StreamReplyContentTypeJSON   string = "application/vnd.zvelo.stream-reply+json"
	StreamsReplyContentType      string = "application/vnd.zvelo.streams-reply"
	StreamsReplyContentTypeJSON  string = "application/vnd.zvelo.streams-reply+json"
	StreamRequestContentType     string = "application/vnd.zvelo.stream-request"
	StreamRequestContentTypeJSON string = "application/vnd.zvelo.stream-request+json"

	SeedResultsContentType     string = "application/vnd.zvelo.seed-results"
	SeedResultsContentTypeJSON string = "application/vnd.zvelo.seed-results+json"
	SeedContentType            string = "application/vnd.zvelo.seed"
	SeedContentTypeJSON        string = "application/vnd.zvelo.seed+json"
)
<commit_msg>remove content types<commit_after><|endoftext|>


<commit_before>// Copyright 2015-2016 trivago GmbH
//
// 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 producer

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/trivago/gollum/core"
	"github.com/trivago/gollum/core/log"
	"github.com/trivago/gollum/shared"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const (
	kinesisCredentialEnv    = "environment"
	kinesisCredentialStatic = "static"
	kinesisCredentialShared = "shared"
	kinesisCredentialNone   = "none"
)

// Kinesis producer plugin
// This producer sends data to an AWS kinesis stream.
// Configuration example
//
//  - "producer.Kinesis":
//    Region: "eu-west-1"
//    Endpoint: "kinesis.eu-west-1.amazonaws.com"
//    CredentialType: "none"
//    CredentialId: ""
//    CredentialToken: ""
//    CredentialSecret: ""
//    CredentialFile: ""
//    CredentialProfile: ""
//    BatchMaxMessages: 500
//    RecordMaxMessages: 1
//    RecordMessageDelimiter: "\n"
//    SendTimeframeSec: 1
//    BatchTimeoutSec: 3
//    StreamMapping:
//      "*" : "default"
//
// KinesisStream defines the stream to read from.
// By default this is set to "default"
//
// Region defines the amazon region of your kinesis stream.
// By default this is set to "eu-west-1".
//
// Endpoint defines the amazon endpoint for your kinesis stream.
// By default this is et to "kinesis.eu-west-1.amazonaws.com"
//
// CredentialType defines the credentials that are to be used when
// connectiong to kensis. This can be one of the following: environment,
// static, shared, none.
// Static enables the parameters CredentialId, CredentialToken and
// CredentialSecret shared enables the parameters CredentialFile and
// CredentialProfile. None will not use any credentials and environment
// will pull the credentials from environmental settings.
// By default this is set to none.
//
// BatchMaxMessages defines the maximum number of messages to send per
// batch. By default this is set to 500.
//
// RecordMaxMessages defines the maximum number of messages to join into
// a kinesis record. By default this is set to 500.
//
// RecordMessageDelimiter defines the string to delimit messages within
// a kinesis record. By default this is set to "\n".
//
// SendTimeframeMs defines the timeframe in milliseconds in which a second
// batch send can be triggered. By default this is set to 1000, i.e. one
// send operation per second.
//
// BatchTimeoutSec defines the number of seconds after which a batch is
// flushed automatically. By default this is set to 3.
//
// StreamMapping defines a translation from gollum stream to kinesis stream
// name. If no mapping is given the gollum stream name is used as kinesis
// stream name.
type Kinesis struct {
	core.ProducerBase
	client            *kinesis.Kinesis
	config            *aws.Config
	streamMap         map[core.MessageStreamID]string
	batch             core.MessageBatch
	recordMaxMessages int
	delimiter         []byte
	flushFrequency    time.Duration
	lastSendTime      time.Time
	sendTimeLimit     time.Duration
	counters          map[string]*int64
	lastMetricUpdate  time.Time
}

const (
	kinesisMetricMessages    = "Kinesis:Messages-"
	kinesisMetricMessagesSec = "Kinesis:MessagesSec-"
)

type streamData struct {
	content            *kinesis.PutRecordsInput
	original           [][]*core.Message
	lastRecordMessages int
}

func init() {
	shared.TypeRegistry.Register(Kinesis{})
}

// Configure initializes this producer with values from a plugin config.
func (prod *Kinesis) Configure(conf core.PluginConfig) error {
	err := prod.ProducerBase.Configure(conf)
	if err != nil {
		return err
	}
	prod.SetStopCallback(prod.close)

	prod.streamMap = conf.GetStreamMap("StreamMapping", "default")
	prod.batch = core.NewMessageBatch(conf.GetInt("BatchMaxMessages", 500))
	prod.recordMaxMessages = conf.GetInt("RecordMaxMessages", 1)
	prod.delimiter = []byte(conf.GetString("RecordMessageDelimiter", "\n"))
	prod.flushFrequency = time.Duration(conf.GetInt("BatchTimeoutSec", 3)) * time.Second
	prod.sendTimeLimit = time.Duration(conf.GetInt("SendTimeframeMs", 1000)) * time.Millisecond
	prod.lastSendTime = time.Now()
	prod.counters = make(map[string]*int64)
	prod.lastMetricUpdate = time.Now()

	if prod.recordMaxMessages < 1 {
		prod.recordMaxMessages = 1
		Log.Warning.Print("RecordMaxMessages was < 1. Defaulting to 1.")
	}

	if prod.recordMaxMessages > 1 && len(prod.delimiter) == 0 {
		prod.delimiter = []byte("\n")
		Log.Warning.Print("RecordMessageDelimiter was empty. Defaulting to \"\\n\".")
	}

	// Config
	prod.config = aws.NewConfig()
	if endpoint := conf.GetString("Endpoint", "kinesis.eu-west-1.amazonaws.com"); endpoint != "" {
		prod.config.WithEndpoint(endpoint)
	}

	if region := conf.GetString("Region", "eu-west-1"); region != "" {
		prod.config.WithRegion(region)
	}

	// Credentials
	credentialType := strings.ToLower(conf.GetString("CredentialType", kinesisCredentialNone))
	switch credentialType {
	case kinesisCredentialEnv:
		prod.config.WithCredentials(credentials.NewEnvCredentials())

	case kinesisCredentialStatic:
		id := conf.GetString("CredentialId", "")
		token := conf.GetString("CredentialToken", "")
		secret := conf.GetString("CredentialSecret", "")
		prod.config.WithCredentials(credentials.NewStaticCredentials(id, secret, token))

	case kinesisCredentialShared:
		filename := conf.GetString("CredentialFile", "")
		profile := conf.GetString("CredentialProfile", "")
		prod.config.WithCredentials(credentials.NewSharedCredentials(filename, profile))

	case kinesisCredentialNone:
		// Nothing

	default:
		return fmt.Errorf("Unknown CredentialType: %s", credentialType)
	}

	for _, streamName := range prod.streamMap {
		shared.Metric.New(kinesisMetricMessages + streamName)
		shared.Metric.New(kinesisMetricMessagesSec + streamName)
		prod.counters[streamName] = new(int64)
	}

	return nil
}

func (prod *Kinesis) bufferMessage(msg core.Message) {
	prod.batch.AppendOrFlush(msg, prod.sendBatch, prod.IsActiveOrStopping, prod.Drop)
}

func (prod *Kinesis) sendBatchOnTimeOut() {
	// Flush if necessary
	if prod.batch.ReachedTimeThreshold(prod.flushFrequency) || prod.batch.ReachedSizeThreshold(prod.batch.Len()/2) {
		prod.sendBatch()
	}

	duration := time.Since(prod.lastMetricUpdate)
	prod.lastMetricUpdate = time.Now()

	for streamName, counter := range prod.counters {
		count := atomic.SwapInt64(counter, 0)

		shared.Metric.Add(kinesisMetricMessages+streamName, count)
		shared.Metric.SetF(kinesisMetricMessagesSec+streamName, float64(count)/duration.Seconds())
	}
}

func (prod *Kinesis) sendBatch() {
	prod.batch.Flush(prod.transformMessages)
}

func (prod *Kinesis) dropMessages(messages []core.Message) {
	for _, msg := range messages {
		prod.Drop(msg)
	}
}

func (prod *Kinesis) transformMessages(messages []core.Message) {
	streamRecords := make(map[core.MessageStreamID]*streamData)

	// Format and sort
	for idx, msg := range messages {
		msgData, streamID := prod.ProducerBase.Format(msg)
		messageHash := fmt.Sprintf("%X-%d", streamID, msg.Sequence)

		// Fetch buffer for this stream
		records, recordsExists := streamRecords[streamID]
		if !recordsExists {
			// Select the correct kinesis stream
			streamName, streamMapped := prod.streamMap[streamID]
			if !streamMapped {
				streamName = core.StreamRegistry.GetStreamName(streamID)
				prod.streamMap[streamID] = streamName

				shared.Metric.New(kinesisMetricMessages + streamName)
				shared.Metric.New(kinesisMetricMessagesSec + streamName)
				prod.counters[streamName] = new(int64)
			}

			// Create buffers for this kinesis stream
			maxLength := len(messages) / prod.recordMaxMessages + 1
			records = &streamData{
				content: &kinesis.PutRecordsInput{
					Records:    make([]*kinesis.PutRecordsRequestEntry, 0, maxLength),
					StreamName: aws.String(streamName),
				},
				original: make([][]*core.Message, 0, maxLength),
				lastRecordMessages: 0,
			}
			streamRecords[streamID] = records
		}

		// Fetch record for this buffer
		var record *kinesis.PutRecordsRequestEntry
		recordExists := len(records.content.Records) > 0
		if !recordExists || records.lastRecordMessages + 1 > prod.recordMaxMessages {
			// Append record to stream
			record = &kinesis.PutRecordsRequestEntry{
				Data:         make([]byte, 0, len(msgData)),
				PartitionKey: aws.String(messageHash),
			}
			records.content.Records = append(records.content.Records, record)
			records.original = append(records.original, make([]*core.Message, 0, prod.recordMaxMessages))
			records.lastRecordMessages = 0
		} else {
			record = records.content.Records[len(records.content.Records)-1]
			record.Data = append(record.Data, prod.delimiter...)
		}

		// Append message to record
		record.Data = append(record.Data, msgData...)
		records.lastRecordMessages += 1
		records.original[len(records.original)-1] = append(records.original[len(records.original)-1], &messages[idx])
	}

	sleepDuration := prod.sendTimeLimit - time.Since(prod.lastSendTime)
	if sleepDuration > 0 {
		time.Sleep(sleepDuration)
	}

	// Send to Kinesis
	for _, records := range streamRecords {
		result, err := prod.client.PutRecords(records.content)
		atomic.AddInt64(prod.counters[*records.content.StreamName], int64(len(records.content.Records)))

		if err != nil {
			// Batch failed, drop all
			Log.Error.Print("Kinesis write error: ", err)
			for _, messages := range records.original {
				for _, msg := range messages {
					prod.Drop(*msg)
				}
			}
		} else {
			// Check each message for errors
			for msgIdx, record := range result.Records {
				if record.ErrorMessage != nil {
					Log.Error.Print("Kinesis message write error: ", *record.ErrorMessage)
					for _, msg := range records.original[msgIdx] {
						prod.Drop(*msg)
					}
				}
			}
		}
	}
}

func (prod *Kinesis) close() {
	defer prod.WorkerDone()
	prod.CloseMessageChannel(prod.bufferMessage)
	prod.batch.Close(prod.transformMessages, prod.GetShutdownTimeout())
}

// Produce writes to stdout or stderr.
func (prod *Kinesis) Produce(workers *sync.WaitGroup) {
	prod.AddMainWorker(workers)

	prod.client = kinesis.New(session.New(prod.config))
	prod.TickerMessageControlLoop(prod.bufferMessage, prod.flushFrequency, prod.sendBatchOnTimeOut)
}
<commit_msg>kinesis producer allow * in StreamMapping<commit_after>// Copyright 2015-2016 trivago GmbH
//
// 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 producer

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/trivago/gollum/core"
	"github.com/trivago/gollum/core/log"
	"github.com/trivago/gollum/shared"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const (
	kinesisCredentialEnv    = "environment"
	kinesisCredentialStatic = "static"
	kinesisCredentialShared = "shared"
	kinesisCredentialNone   = "none"
)

// Kinesis producer plugin
// This producer sends data to an AWS kinesis stream.
// Configuration example
//
//  - "producer.Kinesis":
//    Region: "eu-west-1"
//    Endpoint: "kinesis.eu-west-1.amazonaws.com"
//    CredentialType: "none"
//    CredentialId: ""
//    CredentialToken: ""
//    CredentialSecret: ""
//    CredentialFile: ""
//    CredentialProfile: ""
//    BatchMaxMessages: 500
//    RecordMaxMessages: 1
//    RecordMessageDelimiter: "\n"
//    SendTimeframeSec: 1
//    BatchTimeoutSec: 3
//    StreamMapping:
//      "*" : "default"
//
// KinesisStream defines the stream to read from.
// By default this is set to "default"
//
// Region defines the amazon region of your kinesis stream.
// By default this is set to "eu-west-1".
//
// Endpoint defines the amazon endpoint for your kinesis stream.
// By default this is et to "kinesis.eu-west-1.amazonaws.com"
//
// CredentialType defines the credentials that are to be used when
// connectiong to kensis. This can be one of the following: environment,
// static, shared, none.
// Static enables the parameters CredentialId, CredentialToken and
// CredentialSecret shared enables the parameters CredentialFile and
// CredentialProfile. None will not use any credentials and environment
// will pull the credentials from environmental settings.
// By default this is set to none.
//
// BatchMaxMessages defines the maximum number of messages to send per
// batch. By default this is set to 500.
//
// RecordMaxMessages defines the maximum number of messages to join into
// a kinesis record. By default this is set to 500.
//
// RecordMessageDelimiter defines the string to delimit messages within
// a kinesis record. By default this is set to "\n".
//
// SendTimeframeMs defines the timeframe in milliseconds in which a second
// batch send can be triggered. By default this is set to 1000, i.e. one
// send operation per second.
//
// BatchTimeoutSec defines the number of seconds after which a batch is
// flushed automatically. By default this is set to 3.
//
// StreamMapping defines a translation from gollum stream to kinesis stream
// name. If no mapping is given the gollum stream name is used as kinesis
// stream name.
type Kinesis struct {
	core.ProducerBase
	client            *kinesis.Kinesis
	config            *aws.Config
	streamMap         map[core.MessageStreamID]string
	batch             core.MessageBatch
	recordMaxMessages int
	delimiter         []byte
	flushFrequency    time.Duration
	lastSendTime      time.Time
	sendTimeLimit     time.Duration
	counters          map[string]*int64
	lastMetricUpdate  time.Time
}

const (
	kinesisMetricMessages    = "Kinesis:Messages-"
	kinesisMetricMessagesSec = "Kinesis:MessagesSec-"
)

type streamData struct {
	content            *kinesis.PutRecordsInput
	original           [][]*core.Message
	lastRecordMessages int
}

func init() {
	shared.TypeRegistry.Register(Kinesis{})
}

// Configure initializes this producer with values from a plugin config.
func (prod *Kinesis) Configure(conf core.PluginConfig) error {
	err := prod.ProducerBase.Configure(conf)
	if err != nil {
		return err
	}
	prod.SetStopCallback(prod.close)

	prod.streamMap = conf.GetStreamMap("StreamMapping", "")
	prod.batch = core.NewMessageBatch(conf.GetInt("BatchMaxMessages", 500))
	prod.recordMaxMessages = conf.GetInt("RecordMaxMessages", 1)
	prod.delimiter = []byte(conf.GetString("RecordMessageDelimiter", "\n"))
	prod.flushFrequency = time.Duration(conf.GetInt("BatchTimeoutSec", 3)) * time.Second
	prod.sendTimeLimit = time.Duration(conf.GetInt("SendTimeframeMs", 1000)) * time.Millisecond
	prod.lastSendTime = time.Now()
	prod.counters = make(map[string]*int64)
	prod.lastMetricUpdate = time.Now()

	if prod.recordMaxMessages < 1 {
		prod.recordMaxMessages = 1
		Log.Warning.Print("RecordMaxMessages was < 1. Defaulting to 1.")
	}

	if prod.recordMaxMessages > 1 && len(prod.delimiter) == 0 {
		prod.delimiter = []byte("\n")
		Log.Warning.Print("RecordMessageDelimiter was empty. Defaulting to \"\\n\".")
	}

	// Config
	prod.config = aws.NewConfig()
	if endpoint := conf.GetString("Endpoint", "kinesis.eu-west-1.amazonaws.com"); endpoint != "" {
		prod.config.WithEndpoint(endpoint)
	}

	if region := conf.GetString("Region", "eu-west-1"); region != "" {
		prod.config.WithRegion(region)
	}

	// Credentials
	credentialType := strings.ToLower(conf.GetString("CredentialType", kinesisCredentialNone))
	switch credentialType {
	case kinesisCredentialEnv:
		prod.config.WithCredentials(credentials.NewEnvCredentials())

	case kinesisCredentialStatic:
		id := conf.GetString("CredentialId", "")
		token := conf.GetString("CredentialToken", "")
		secret := conf.GetString("CredentialSecret", "")
		prod.config.WithCredentials(credentials.NewStaticCredentials(id, secret, token))

	case kinesisCredentialShared:
		filename := conf.GetString("CredentialFile", "")
		profile := conf.GetString("CredentialProfile", "")
		prod.config.WithCredentials(credentials.NewSharedCredentials(filename, profile))

	case kinesisCredentialNone:
		// Nothing

	default:
		return fmt.Errorf("Unknown CredentialType: %s", credentialType)
	}

	for _, streamName := range prod.streamMap {
		shared.Metric.New(kinesisMetricMessages + streamName)
		shared.Metric.New(kinesisMetricMessagesSec + streamName)
		prod.counters[streamName] = new(int64)
	}

	return nil
}

func (prod *Kinesis) bufferMessage(msg core.Message) {
	prod.batch.AppendOrFlush(msg, prod.sendBatch, prod.IsActiveOrStopping, prod.Drop)
}

func (prod *Kinesis) sendBatchOnTimeOut() {
	// Flush if necessary
	if prod.batch.ReachedTimeThreshold(prod.flushFrequency) || prod.batch.ReachedSizeThreshold(prod.batch.Len()/2) {
		prod.sendBatch()
	}

	duration := time.Since(prod.lastMetricUpdate)
	prod.lastMetricUpdate = time.Now()

	for streamName, counter := range prod.counters {
		count := atomic.SwapInt64(counter, 0)

		shared.Metric.Add(kinesisMetricMessages+streamName, count)
		shared.Metric.SetF(kinesisMetricMessagesSec+streamName, float64(count)/duration.Seconds())
	}
}

func (prod *Kinesis) sendBatch() {
	prod.batch.Flush(prod.transformMessages)
}

func (prod *Kinesis) dropMessages(messages []core.Message) {
	for _, msg := range messages {
		prod.Drop(msg)
	}
}

func (prod *Kinesis) transformMessages(messages []core.Message) {
	streamRecords := make(map[core.MessageStreamID]*streamData)

	// Format and sort
	for idx, msg := range messages {
		msgData, streamID := prod.ProducerBase.Format(msg)
		messageHash := fmt.Sprintf("%X-%d", streamID, msg.Sequence)

		// Fetch buffer for this stream
		records, recordsExists := streamRecords[streamID]
		if !recordsExists {
			// Select the correct kinesis stream
			streamName, streamMapped := prod.streamMap[streamID]
			if !streamMapped {
				streamName, streamMapped = prod.streamMap[core.WildcardStreamID]
				if !streamMapped {
					streamName = core.StreamRegistry.GetStreamName(streamID)
					prod.streamMap[streamID] = streamName

					shared.Metric.New(kinesisMetricMessages + streamName)
					shared.Metric.New(kinesisMetricMessagesSec + streamName)
					prod.counters[streamName] = new(int64)
				}
			}

			// Create buffers for this kinesis stream
			maxLength := len(messages) / prod.recordMaxMessages + 1
			records = &streamData{
				content: &kinesis.PutRecordsInput{
					Records:    make([]*kinesis.PutRecordsRequestEntry, 0, maxLength),
					StreamName: aws.String(streamName),
				},
				original: make([][]*core.Message, 0, maxLength),
				lastRecordMessages: 0,
			}
			streamRecords[streamID] = records
		}

		// Fetch record for this buffer
		var record *kinesis.PutRecordsRequestEntry
		recordExists := len(records.content.Records) > 0
		if !recordExists || records.lastRecordMessages + 1 > prod.recordMaxMessages {
			// Append record to stream
			record = &kinesis.PutRecordsRequestEntry{
				Data:         make([]byte, 0, len(msgData)),
				PartitionKey: aws.String(messageHash),
			}
			records.content.Records = append(records.content.Records, record)
			records.original = append(records.original, make([]*core.Message, 0, prod.recordMaxMessages))
			records.lastRecordMessages = 0
		} else {
			record = records.content.Records[len(records.content.Records)-1]
			record.Data = append(record.Data, prod.delimiter...)
		}

		// Append message to record
		record.Data = append(record.Data, msgData...)
		records.lastRecordMessages += 1
		records.original[len(records.original)-1] = append(records.original[len(records.original)-1], &messages[idx])
	}

	sleepDuration := prod.sendTimeLimit - time.Since(prod.lastSendTime)
	if sleepDuration > 0 {
		time.Sleep(sleepDuration)
	}

	// Send to Kinesis
	for _, records := range streamRecords {
		result, err := prod.client.PutRecords(records.content)
		atomic.AddInt64(prod.counters[*records.content.StreamName], int64(len(records.content.Records)))

		if err != nil {
			// Batch failed, drop all
			Log.Error.Print("Kinesis write error: ", err)
			for _, messages := range records.original {
				for _, msg := range messages {
					prod.Drop(*msg)
				}
			}
		} else {
			// Check each message for errors
			for msgIdx, record := range result.Records {
				if record.ErrorMessage != nil {
					Log.Error.Print("Kinesis message write error: ", *record.ErrorMessage)
					for _, msg := range records.original[msgIdx] {
						prod.Drop(*msg)
					}
				}
			}
		}
	}
}

func (prod *Kinesis) close() {
	defer prod.WorkerDone()
	prod.CloseMessageChannel(prod.bufferMessage)
	prod.batch.Close(prod.transformMessages, prod.GetShutdownTimeout())
}

// Produce writes to stdout or stderr.
func (prod *Kinesis) Produce(workers *sync.WaitGroup) {
	prod.AddMainWorker(workers)

	prod.client = kinesis.New(session.New(prod.config))
	prod.TickerMessageControlLoop(prod.bufferMessage, prod.flushFrequency, prod.sendBatchOnTimeOut)
}
<|endoftext|>


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

package abelana

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"appengine"
	"appengine/datastore"
	"appengine/socket"

	"github.com/garyburd/redigo/redis"
)

var server string

// Note - I looked at adapting both Gary Burd's pool system and Vites pools to AppEngine, but ran
// out of time as there were too many dependencies.

func redisInit() {
	if appengine.IsDevAppServer() {
		server = redisExt
	} else {
		server = redisExt
	}
}

// iNowFollow is Called when we need to add a folower
func iNowFollow(cx appengine.Context, userID, followerID string) {

}

// addPhoto is called to add a photo.
func addPhoto(c appengine.Context, filename string) {
	err := func() error {
		s := strings.Split(filename, ".")
		if len(s) != 3 {
			return fmt.Errorf("bad filename format %v", filename)
		}
		// s[2] is the file extension
		userID, photoID, superID := s[0], s[1], s[0]+"."+s[1]
		u, err := findUser(c, userID)
		if err != nil {
			return fmt.Errorf("unable to find user %v %v", userID, err)
		}

		p := &Photo{photoID, time.Now().UTC().Unix()}

		socket, err := socket.Dial(c, "tcp", server)
		if err != nil {
			return fmt.Errorf("Dial %v", err)
		}
		conn := redis.NewConn(socket, 0, 0) // TODO 0 TO's for now
		defer conn.Close()

		// Set Date
		set, err := redis.Int(conn.Do("HSETNX", "IM:"+superID, "date", p.Date))
		if (err != nil && err != redis.ErrNil) || set == 0 {
			// block duplicate requests
			return fmt.Errorf("duplicate %v %v", err, set)
		}

		// TODO: Consider that at some point these should be done in batches of 100 or so.

		// Add to each follower's list
		for _, personID := range u.Persons {
			conn.Send("LPUSH", "TL:"+personID, photoID)
		}
		conn.Flush()

		// Check the result and adjust list if nescessary.
		for _, personID := range u.Persons {
			v, err := redis.Int(conn.Receive())
			if err != nil && err != redis.ErrNil {
				c.Errorf("AddPhoto: TL:%v %v", personID, err)
				continue
			}
			if v > 2000 {
				_, err := conn.Do("RPOP", "TL:"+personID)
				if err != nil {
					c.Errorf("AddPhoto RPOP %v", err)
				}
			}
		}
		k1 := datastore.NewKey(c, "User", s[0], 0, nil)
		k2 := datastore.NewKey(c, "Photo", s[1], 0, k1)
		_, err = datastore.Put(c, k2, p)
		if err != nil {
			return fmt.Errorf("datastore %v", err)
		}
		return nil
	}()
	if err != nil {
		c.Errorf("AddPhoto: %v", err)
	}
}

// getTimeline
func getTimeline(cx appengine.Context, userID, lastid string) ([]TLEntry, error) {
	var item string
	timeline := []TLEntry{}

	hc, err := socket.Dial(cx, "tcp", server)
	if err != nil {
		cx.Errorf("GetTimeLine Dial %v", err)
		return nil, err
	}
	defer hc.Close()
	conn := redis.NewConn(hc, 0, 0) // TODO 0 TO's for now

	list, err := redis.Strings(conn.Do("LRANGE", "TL:"+userID, 0, -1))
	if err != nil && err != redis.ErrNil {
		cx.Errorf("GetTimeLine %v", err)
	}
	idx := 0
	find := func(vs []int, x int) (int, bool) {
		for i, v := range vs {
			if v == x {
				return i, true
			}
		}
		return 0, false
	}

	if lastid != "0" {
		for i := 0; i < len(list) && list[i] != lastid; i++ {

		}
		for i, item := range list {
			if item == lastid {
				ix = i
				break
			}
		}
	}
	timeline = make([]TLEntry, 0, timelineBatchSize)
	for i := 0; i < timelineBatchSize && i+ix < len(list); i++ {
		photoID := list[ix+i]

		v, err := redis.Strings(conn.Do("HMGET", "IM:"+photoID, "date", userID, "flag"))
		if err != nil && err != redis.ErrNil {
			cx.Errorf("GetTimeLine HMGET %v", err)
		}
		if v[2] != "" {
			flags, err := strconv.Atoi(v[2])
			if err == nil && flags > 1 {
				continue // skip flag'd images
			}
		}
		likes, err := redis.Int(conn.Do("HLEN", "IM:"+photoID))
		if err != nil && err != redis.ErrNil {
			cx.Errorf("GetTimeLine HLEN %v", err)
		}
		s := strings.Split(photoID, ".")
		dn, err := redis.String(conn.Do("HGET", "HT:"+s[0], "dn"))
		if err != nil && err != redis.ErrNil {
			cx.Errorf("GetTimeLine HLEN %v", err)
		}
		dt, err := strconv.ParseInt(v[0], 10, 64)
		te := TLEntry{dt, s[0], dn, photoID, likes - 1, v[1] == "1"}
		timeline = append(timeline, te)
	}
	return timeline, nil
}

// addUser adds the user to redis
func addUser(cx appengine.Context, userID, displayName string) {
	hc, err := socket.Dial(cx, "tcp", server)
	if err != nil {
		cx.Errorf("addUser Dial %v", err)
	}
	defer hc.Close()
	conn := redis.NewConn(hc, 0, 0) // TODO 0 TO's for now

	// See if we have done this already, block others.
	_, err = conn.Do("HSET", "HT:"+userID, "dn", displayName)
	if err != nil {
		cx.Errorf("addUser Exists %v", err)
	}
}

// like the user on redis
func like(cx appengine.Context, userID, photoID string) {
	hc, err := socket.Dial(cx, "tcp", server)
	if err != nil {
		cx.Errorf("like Dial %v", err)
		return
	}
	defer hc.Close()
	conn := redis.NewConn(hc, 0, 0) // TODO 0 TO's for now
	_, err = redis.Int(conn.Do("HSET", "IM:"+photoID, userID, "1"))
	if err != nil && err != redis.ErrNil {
		cx.Errorf("like %v", err)
	}
}

// unlike
func unlike(cx appengine.Context, userID, photoID string) {
	hc, err := socket.Dial(cx, "tcp", server)
	if err != nil {
		cx.Errorf("unlike Dial %v", err)
		return
	}
	defer hc.Close()
	conn := redis.NewConn(hc, 0, 0) // TODO 0 TO's for now
	_, err = redis.Int(conn.Do("HDEL", "IM:"+photoID, userID))
	if err != nil && err != redis.ErrNil {
		cx.Errorf("unlike %v", err)
	}
}

// flag will tell us that things may not be quite right with this image.
func flag(cx appengine.Context, userID, photoID string) {
	hc, err := socket.Dial(cx, "tcp", server)
	if err != nil {
		cx.Errorf("unlike Dial %v", err)
		return
	}
	defer hc.Close()
	conn := redis.NewConn(hc, 0, 0) // TODO 0 TO's for now
	_, err = redis.Int(conn.Do("HINCRBY", "IM:"+photoID, "flag", 1))
	if err != nil && err != redis.ErrNil {
		cx.Errorf("unlike %v", err)
	}
}
<commit_msg>back to older version for easier merge<commit_after>// Copyright 2014 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package abelana

import (
	"strconv"
	"strings"
	"time"

	"appengine"
	"appengine/datastore"
	"appengine/socket"

	"github.com/garyburd/redigo/redis"
)

var server string

// Note - I looked at adapting both Gary Burd's pool system and Vites pools to AppEngine, but ran
// out of time as there were too many dependencies.

func redisInit() {
	if appengine.IsDevAppServer() {
		server = redisExt
	} else {
		server = redisExt
	}
}

// iNowFollow is Called when we need to add a folower
func iNowFollow(cx appengine.Context, userID, followerID string) {

}

// addPhoto is called to add a photo.
func addPhoto(cx appengine.Context, superID string) {

	s := strings.Split(superID, ".")
	if len(s) != 3 {
		cx.Errorf("AddPhoto -- bad superID %v", superID)
		return
	}
	u, err := findUser(cx, s[0])
	if err != nil {
		cx.Errorf("AddPhoto unable to find user %v %v", superID, err)
		return
	}

	p := &Photo{s[1], time.Now().UTC().Unix()}
	hc, err := socket.Dial(cx, "tcp", server)
	if err != nil {
		cx.Errorf("AddPhoto Dial %v", err)
		return
	}
	defer hc.Close()
	conn := redis.NewConn(hc, 0, 0) // TODO 0 TO's for now

	photoID := s[0] + "." + s[1]
	set, err := redis.Int(conn.Do("HSETNX", "IM:"+photoID, "date", p.Date)) // Set Date
	if (err != nil && err != redis.ErrNil) || set == 0 {
		cx.Errorf("AddPhoto D duplicate %v %v", err, set)
		return // block duplicate requests
	}

	// Consider that at some point these should be done in batches of 100 or so. TODO

	// Add to each follower's list
	for _, personID := range u.Persons {
		conn.Send("LPUSH", "TL:"+personID, photoID)
	}
	conn.Flush()

	// Check the result and adjust list if nescessary.
	for _, personID := range u.Persons {
		v, err := redis.Int(conn.Receive())
		if err != nil && err != redis.ErrNil {
			cx.Errorf("AddPhoto for %v %v", "TL:"+personID, err)
		} else {
			if v > 2000 {
				_, err := conn.Do("RPOP", "TL:"+personID)
				if err != nil {
					cx.Errorf("AddPhoto RPOP %v", err)
				}
			}
		}
	}
	k1 := datastore.NewKey(cx, "User", s[0], 0, nil)
	k2 := datastore.NewKey(cx, "Photo", s[1], 0, k1)
	_, err = datastore.Put(cx, k2, p)
	if err != nil {
		cx.Errorf("AddPhoto datastore %v", err)
	}
}

// getTimeline
func getTimeline(cx appengine.Context, userID, lastid string) ([]TLEntry, error) {
	var item string
	timeline := []TLEntry{}

	hc, err := socket.Dial(cx, "tcp", server)
	if err != nil {
		cx.Errorf("GetTimeLine Dial %v", err)
		return nil, err
	}
	defer hc.Close()
	conn := redis.NewConn(hc, 0, 0) // TODO 0 TO's for now

	list, err := redis.Strings(conn.Do("LRANGE", "TL:"+userID, 0, -1))
	if err != nil && err != redis.ErrNil {
		cx.Errorf("GetTimeLine %v", err)
	}
	ix := 0
	if lastid != "0" {
		for ix, item = range list {
			if item == lastid {
				break
			}
		}
	}
	timeline = make([]TLEntry, 0, timelineBatchSize)
	for i := 0; i < timelineBatchSize && i+ix < len(list); i++ {
		photoID := list[ix+i]

		v, err := redis.Strings(conn.Do("HMGET", "IM:"+photoID, "date", userID, "flag"))
		if err != nil && err != redis.ErrNil {
			cx.Errorf("GetTimeLine HMGET %v", err)
		}
		if v[2] != "" {
			flags, err := strconv.Atoi(v[2])
			if err == nil && flags > 1 {
				continue // skip flag'd images
			}
		}
		likes, err := redis.Int(conn.Do("HLEN", "IM:"+photoID))
		if err != nil && err != redis.ErrNil {
			cx.Errorf("GetTimeLine HLEN %v", err)
		}
		s := strings.Split(photoID, ".")
		dn, err := redis.String(conn.Do("HGET", "HT:"+s[0], "dn"))
		if err != nil && err != redis.ErrNil {
			cx.Errorf("GetTimeLine HLEN %v", err)
		}
		dt, err := strconv.ParseInt(v[0], 10, 64)
		te := TLEntry{dt, s[0], dn, photoID, likes - 1, v[1] == "1"}
		timeline = append(timeline, te)
	}
	return timeline, nil
}

// addUser adds the user to redis
func addUser(cx appengine.Context, userID, displayName string) {
	hc, err := socket.Dial(cx, "tcp", server)
	if err != nil {
		cx.Errorf("addUser Dial %v", err)
	}
	defer hc.Close()
	conn := redis.NewConn(hc, 0, 0) // TODO 0 TO's for now

	// See if we have done this already, block others.
	_, err = conn.Do("HSET", "HT:"+userID, "dn", displayName)
	if err != nil {
		cx.Errorf("addUser Exists %v", err)
	}
}

// like the user on redis
func like(cx appengine.Context, userID, photoID string) {
	hc, err := socket.Dial(cx, "tcp", server)
	if err != nil {
		cx.Errorf("like Dial %v", err)
		return
	}
	defer hc.Close()
	conn := redis.NewConn(hc, 0, 0) // TODO 0 TO's for now
	_, err = redis.Int(conn.Do("HSET", "IM:"+photoID, userID, "1"))
	if err != nil && err != redis.ErrNil {
		cx.Errorf("like %v", err)
	}
}

// unlike
func unlike(cx appengine.Context, userID, photoID string) {
	hc, err := socket.Dial(cx, "tcp", server)
	if err != nil {
		cx.Errorf("unlike Dial %v", err)
		return
	}
	defer hc.Close()
	conn := redis.NewConn(hc, 0, 0) // TODO 0 TO's for now
	_, err = redis.Int(conn.Do("HDEL", "IM:"+photoID, userID))
	if err != nil && err != redis.ErrNil {
		cx.Errorf("unlike %v", err)
	}
}

// flag will tell us that things may not be quite right with this image.
func flag(cx appengine.Context, userID, photoID string) {
	hc, err := socket.Dial(cx, "tcp", server)
	if err != nil {
		cx.Errorf("unlike Dial %v", err)
		return
	}
	defer hc.Close()
	conn := redis.NewConn(hc, 0, 0) // TODO 0 TO's for now
	_, err = redis.Int(conn.Do("HINCRBY", "IM:"+photoID, "flag", 1))
	if err != nil && err != redis.ErrNil {
		cx.Errorf("unlike %v", err)
	}
}
<|endoftext|>


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

package casbin

import (
	"log"
	"sync"
	"time"
)

type SyncedEnforcer struct {
	*Enforcer
	m sync.RWMutex
	autoLoad bool
}

// NewEnforcer creates a synchronized enforcer via file or DB.
func NewSyncedEnforcer(params ...interface{}) *SyncedEnforcer {
	e := &SyncedEnforcer{}
	e.Enforcer = NewEnforcer(params...)
	e.autoLoad = false
	return e
}

func (e *SyncedEnforcer) PeriodicallyLoadPolicy(d time.Duration) {
	if e.autoLoad {
		time.Sleep(30 * time.Millisecond)
		e.LoadPolicy()
	}
}

func (e *SyncedEnforcer) StartAutoLoadPolicy(d time.Duration) {
	e.autoLoad = true
	go func() {
		n := 1
		log.Print("Start automatically load policy")
		for {
			if !e.autoLoad {
				log.Print("Stop automatically load policy")
				break
			}

			e.m.Lock()
			e.LoadPolicy()
			e.m.Unlock()
			// Uncomment this line to see when the policy is loaded.
			// log.Print("Load policy for time: ", n)
			n ++
			time.Sleep(d)
		}
	} ()
}

func (e *SyncedEnforcer) StopAutoLoadPolicy() {
	e.autoLoad = false
}

// Enforce decides whether a "subject" can access a "object" with the operation "action", input parameters are usually: (sub, obj, act).
func (e *SyncedEnforcer) Enforce(rvals ...interface{}) bool {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.Enforce(rvals...)
}

// GetAllSubjects gets the list of subjects that show up in the current policy.
func (e *SyncedEnforcer) GetAllSubjects() []string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetAllSubjects()
}

// GetAllObjects gets the list of objects that show up in the current policy.
func (e *SyncedEnforcer) GetAllObjects() []string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetAllObjects()
}

// GetAllActions gets the list of actions that show up in the current policy.
func (e *SyncedEnforcer) GetAllActions() []string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetAllActions()
}

// GetAllRoles gets the list of roles that show up in the current policy.
func (e *SyncedEnforcer) GetAllRoles() []string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetAllRoles()
}

// GetPolicy gets all the authorization rules in the policy.
func (e *SyncedEnforcer) GetPolicy() [][]string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetPolicy()
}

// GetFilteredPolicy gets all the authorization rules in the policy, field filters can be specified.
func (e *SyncedEnforcer) GetFilteredPolicy(fieldIndex int, fieldValues ...string) [][]string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetFilteredPolicy(fieldIndex, fieldValues...)
}

// GetGroupingPolicy gets all the role inheritance rules in the policy.
func (e *SyncedEnforcer) GetGroupingPolicy() [][]string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetGroupingPolicy()
}

// GetFilteredGroupingPolicy gets all the role inheritance rules in the policy, field filters can be specified.
func (e *SyncedEnforcer) GetFilteredGroupingPolicy(fieldIndex int, fieldValues ...string) [][]string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetFilteredGroupingPolicy(fieldIndex, fieldValues...)
}

// HasPolicy determines whether an authorization rule exists.
func (e *SyncedEnforcer) HasPolicy(params ...interface{}) bool {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.HasPolicy(params...)
}

// AddPolicy adds an authorization rule to the current policy.
// If the rule already exists, the function returns false and the rule will not be added.
// Otherwise the function returns true by adding the new rule.
func (e *SyncedEnforcer) AddPolicy(params ...interface{}) bool {
	e.m.Lock()
	defer e.m.Unlock()
	return e.Enforcer.AddPolicy(params...)
}

// RemovePolicy removes an authorization rule from the current policy.
func (e *SyncedEnforcer) RemovePolicy(params ...interface{}) bool {
	e.m.Lock()
	defer e.m.Unlock()
	return e.Enforcer.RemovePolicy(params...)
}

// RemoveFilteredPolicy removes an authorization rule from the current policy, field filters can be specified.
func (e *SyncedEnforcer) RemoveFilteredPolicy(fieldIndex int, fieldValues ...string) bool {
	e.m.Lock()
	defer e.m.Unlock()
	return e.Enforcer.RemoveFilteredPolicy(fieldIndex, fieldValues...)
}

// HasGroupingPolicy determines whether a role inheritance rule exists.
func (e *SyncedEnforcer) HasGroupingPolicy(params ...interface{}) bool {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.HasGroupingPolicy(params...)
}

// AddGroupingPolicy adds a role inheritance rule to the current policy.
// If the rule already exists, the function returns false and the rule will not be added.
// Otherwise the function returns true by adding the new rule.
func (e *SyncedEnforcer) AddGroupingPolicy(params ...interface{}) bool {
	e.m.Lock()
	defer e.m.Unlock()
	return e.Enforcer.AddGroupingPolicy(params...)
}

// RemoveGroupingPolicy removes a role inheritance rule from the current policy.
func (e *SyncedEnforcer) RemoveGroupingPolicy(params ...interface{}) bool {
	e.m.Lock()
	defer e.m.Unlock()
	return e.Enforcer.RemoveGroupingPolicy(params...)
}

// RemoveFilteredGroupingPolicy removes a role inheritance rule from the current policy, field filters can be specified.
func (e *SyncedEnforcer) RemoveFilteredGroupingPolicy(fieldIndex int, fieldValues ...string) bool {
	e.m.Lock()
	defer e.m.Unlock()
	return e.Enforcer.RemoveFilteredGroupingPolicy(fieldIndex, fieldValues...)
}
<commit_msg>Add LoadPolicy() and SavePolicy() to SyncedEnforcer.<commit_after>// Copyright 2017 The casbin Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package casbin

import (
	"log"
	"sync"
	"time"
)

type SyncedEnforcer struct {
	*Enforcer
	m sync.RWMutex
	autoLoad bool
}

// NewEnforcer creates a synchronized enforcer via file or DB.
func NewSyncedEnforcer(params ...interface{}) *SyncedEnforcer {
	e := &SyncedEnforcer{}
	e.Enforcer = NewEnforcer(params...)
	e.autoLoad = false
	return e
}

func (e *SyncedEnforcer) PeriodicallyLoadPolicy(d time.Duration) {
	if e.autoLoad {
		time.Sleep(30 * time.Millisecond)
		e.LoadPolicy()
	}
}

func (e *SyncedEnforcer) StartAutoLoadPolicy(d time.Duration) {
	e.autoLoad = true
	go func() {
		n := 1
		log.Print("Start automatically load policy")
		for {
			if !e.autoLoad {
				log.Print("Stop automatically load policy")
				break
			}

			e.LoadPolicy()
			// Uncomment this line to see when the policy is loaded.
			// log.Print("Load policy for time: ", n)
			n ++
			time.Sleep(d)
		}
	} ()
}

func (e *SyncedEnforcer) StopAutoLoadPolicy() {
	e.autoLoad = false
}

// LoadPolicy reloads the policy from file/database.
func (e *SyncedEnforcer) LoadPolicy() error {
	e.m.Lock()
	defer e.m.Unlock()
	return e.Enforcer.LoadPolicy()
}

// SavePolicy saves the current policy (usually after changed with Casbin API) back to file/database.
func (e *SyncedEnforcer) SavePolicy() error {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.SavePolicy()
}

// Enforce decides whether a "subject" can access a "object" with the operation "action", input parameters are usually: (sub, obj, act).
func (e *SyncedEnforcer) Enforce(rvals ...interface{}) bool {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.Enforce(rvals...)
}

// GetAllSubjects gets the list of subjects that show up in the current policy.
func (e *SyncedEnforcer) GetAllSubjects() []string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetAllSubjects()
}

// GetAllObjects gets the list of objects that show up in the current policy.
func (e *SyncedEnforcer) GetAllObjects() []string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetAllObjects()
}

// GetAllActions gets the list of actions that show up in the current policy.
func (e *SyncedEnforcer) GetAllActions() []string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetAllActions()
}

// GetAllRoles gets the list of roles that show up in the current policy.
func (e *SyncedEnforcer) GetAllRoles() []string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetAllRoles()
}

// GetPolicy gets all the authorization rules in the policy.
func (e *SyncedEnforcer) GetPolicy() [][]string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetPolicy()
}

// GetFilteredPolicy gets all the authorization rules in the policy, field filters can be specified.
func (e *SyncedEnforcer) GetFilteredPolicy(fieldIndex int, fieldValues ...string) [][]string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetFilteredPolicy(fieldIndex, fieldValues...)
}

// GetGroupingPolicy gets all the role inheritance rules in the policy.
func (e *SyncedEnforcer) GetGroupingPolicy() [][]string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetGroupingPolicy()
}

// GetFilteredGroupingPolicy gets all the role inheritance rules in the policy, field filters can be specified.
func (e *SyncedEnforcer) GetFilteredGroupingPolicy(fieldIndex int, fieldValues ...string) [][]string {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.GetFilteredGroupingPolicy(fieldIndex, fieldValues...)
}

// HasPolicy determines whether an authorization rule exists.
func (e *SyncedEnforcer) HasPolicy(params ...interface{}) bool {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.HasPolicy(params...)
}

// AddPolicy adds an authorization rule to the current policy.
// If the rule already exists, the function returns false and the rule will not be added.
// Otherwise the function returns true by adding the new rule.
func (e *SyncedEnforcer) AddPolicy(params ...interface{}) bool {
	e.m.Lock()
	defer e.m.Unlock()
	return e.Enforcer.AddPolicy(params...)
}

// RemovePolicy removes an authorization rule from the current policy.
func (e *SyncedEnforcer) RemovePolicy(params ...interface{}) bool {
	e.m.Lock()
	defer e.m.Unlock()
	return e.Enforcer.RemovePolicy(params...)
}

// RemoveFilteredPolicy removes an authorization rule from the current policy, field filters can be specified.
func (e *SyncedEnforcer) RemoveFilteredPolicy(fieldIndex int, fieldValues ...string) bool {
	e.m.Lock()
	defer e.m.Unlock()
	return e.Enforcer.RemoveFilteredPolicy(fieldIndex, fieldValues...)
}

// HasGroupingPolicy determines whether a role inheritance rule exists.
func (e *SyncedEnforcer) HasGroupingPolicy(params ...interface{}) bool {
	e.m.RLock()
	defer e.m.RUnlock()
	return e.Enforcer.HasGroupingPolicy(params...)
}

// AddGroupingPolicy adds a role inheritance rule to the current policy.
// If the rule already exists, the function returns false and the rule will not be added.
// Otherwise the function returns true by adding the new rule.
func (e *SyncedEnforcer) AddGroupingPolicy(params ...interface{}) bool {
	e.m.Lock()
	defer e.m.Unlock()
	return e.Enforcer.AddGroupingPolicy(params...)
}

// RemoveGroupingPolicy removes a role inheritance rule from the current policy.
func (e *SyncedEnforcer) RemoveGroupingPolicy(params ...interface{}) bool {
	e.m.Lock()
	defer e.m.Unlock()
	return e.Enforcer.RemoveGroupingPolicy(params...)
}

// RemoveFilteredGroupingPolicy removes a role inheritance rule from the current policy, field filters can be specified.
func (e *SyncedEnforcer) RemoveFilteredGroupingPolicy(fieldIndex int, fieldValues ...string) bool {
	e.m.Lock()
	defer e.m.Unlock()
	return e.Enforcer.RemoveFilteredGroupingPolicy(fieldIndex, fieldValues...)
}
<|endoftext|>


<commit_before>package ITC

import (
	"fmt"
	"strings"
)

// An Event is a recursively defined binary interval tree
// with
type Event struct {
	n      int    // Value of this node
	e1, e2 *Event // left child, right child
}

// An Id is a recursively defined binary interval tree
// which can be divded or merged with other Ids.
type Id struct {
	n      int // Value of this node
	i1, i2 *Id // left child, right child
}

// A stamp is a causal event tracker. It is composed
// of an Id (the who) and an Event (the what).
// It does not encode any payload data, merely a causal
// history which may be attached to a message.
type Stamp struct {
	event *Event
	id    *Id
}

// ================= FUNCTIONS ===============

// Convert the id into a compact string for printing
func (id *Id) String() string {
	switch {
	case id.n == -1:
		l, r := "nil", "nil"
		if id.i1 != nil {
			l = id.i1.String()
		}
		if id.i2 != nil {
			r = id.i2.String()
		}
		return fmt.Sprintf("(%s, %s)", l, r)
	default:
		return fmt.Sprintf("%d", id.n)
	}
}

// Reverse the String conversion
func stringToId(s string) *Id {
	return sReaderToId(strings.NewReader(s))
}

// Recursive helper function that consumes a reader
func sReaderToId(s *strings.Reader) (id *Id) {
	id = &Id{n: -1}
Loop:
	for {
		ch, _, err := s.ReadRune()
		switch {
		case err != nil:
			break Loop
		case ch == ')':
			return
		case ch == '(' || ch == ',':
			nId := sReaderToId(s)
			if ch == '(' {
				id.i1 = nId
			} else {
				id.i2 = nId
			}
		case ch == '1' || ch == '0':
			return &Id{n: int(ch) - '0'}
		default:
		}
	}
	return
}
<commit_msg>Added comments<commit_after>package ITC

import (
	"fmt"
	"strings"
)

// An Event is a recursively defined binary interval tree
// with
type Event struct {
	n      int    // Value of this node
	e1, e2 *Event // left child, right child
}

// An Id is a recursively defined binary interval tree
// which can be divded or merged with other Ids.
type Id struct {
	n      int // Value of this node
	i1, i2 *Id // left child, right child
}

// A stamp is a causal event tracker. It is composed
// of an Id (the who) and an Event (the what).
// It does not encode any payload data, merely a causal
// history which may be attached to a message.
type Stamp struct {
	event *Event
	id    *Id
}

// ================= FUNCTIONS ===============

// Convert the id into a compact string for printing
func (id *Id) String() string {
	switch {
	case id.n == -1:
		l, r := "nil", "nil"
		if id.i1 != nil {
			l = id.i1.String()
		}
		if id.i2 != nil {
			r = id.i2.String()
		}
		return fmt.Sprintf("(%s, %s)", l, r)
	default:
		return fmt.Sprintf("%d", id.n)
	}
}

// Reverse the String conversion
func stringToId(s string) *Id {
	return sReaderToId(strings.NewReader(s))
}

// Recursive helper function that consumes a reader
func sReaderToId(s *strings.Reader) (id *Id) {
	id = &Id{n: -1}
	for { // Iterate
		ch, _, err := s.ReadRune()
		switch {
		// right paren or end of string ends the id
		case err != nil, ch == ')':
			return
		case ch == '(' || ch == ',':
			// left paren begins a new left tree,
			// comma begins a new right tree
			nId := sReaderToId(s)
			if ch == '(' {
				id.i1 = nId
			} else {
				id.i2 = nId
			}
		// Atoms are always nodes of their own
		case ch == '1' || ch == '0':
			return &Id{n: int(ch) - '0'}
		default:
		}
	}
	return
}
<|endoftext|>


<commit_before>package api

import (
	"context"
	"net/http"
	"time"

	"github.com/gorilla/mux"
	"github.com/rockbears/log"

	"github.com/ovh/cds/engine/api/authentication"
	"github.com/ovh/cds/engine/api/group"
	"github.com/ovh/cds/engine/api/user"
	"github.com/ovh/cds/engine/service"
	"github.com/ovh/cds/sdk"
)

func (api *API) getAuthDriversHandler() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		drivers := make(sdk.AuthDriverManifests, 0, len(api.AuthenticationDrivers))

		for _, d := range api.AuthenticationDrivers {
			drivers = append(drivers, d.GetManifest())
		}

		countAdmins, err := user.CountAdmin(api.mustDB())
		if err != nil {
			return err
		}

		var response = sdk.AuthDriverResponse{
			IsFirstConnection: countAdmins == 0,
			Drivers:           drivers,
		}

		return service.WriteJSON(w, response, http.StatusOK)
	}
}

func (api *API) getAuthScopesHandler() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		return service.WriteJSON(w, api.Router.scopeDetails, http.StatusOK)
	}
}

func (api *API) getAuthAskSigninHandler() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		vars := mux.Vars(r)

		// Extract consumer type from request, is invalid or not in api drivers list return an error
		consumerType := sdk.AuthConsumerType(vars["consumerType"])
		if !consumerType.IsValid() {
			return sdk.WithStack(sdk.ErrNotFound)
		}
		driver, ok := api.AuthenticationDrivers[consumerType]
		if !ok {
			return sdk.WithStack(sdk.ErrNotFound)
		}

		driverRedirect, ok := driver.(sdk.AuthDriverWithRedirect)
		if !ok {
			return nil
		}

		countAdmins, err := user.CountAdmin(api.mustDB())
		if err != nil {
			return err
		}

		var signinState = sdk.AuthSigninConsumerToken{
			RequireMFA:        QueryBool(r, "require_mfa"),
			RedirectURI:       QueryString(r, "redirect_uri"),
			IsFirstConnection: countAdmins == 0,
		}
		// Get the origin from request if set
		signinState.Origin = QueryString(r, "origin")

		if signinState.Origin != "" && !(signinState.Origin == "cdsctl" || signinState.Origin == "ui") {
			return sdk.NewErrorFrom(sdk.ErrWrongRequest, "invalid given origin value")
		}

		// Redirect to the right signin page depending on the consumer type
		redirect, err := driverRedirect.GetSigninURI(signinState)
		if err != nil {
			return err
		}

		return service.WriteJSON(w, redirect, http.StatusOK)
	}
}

func (api *API) postAuthSigninHandler() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		vars := mux.Vars(r)

		// Extract consumer type from request, is invalid or not in api drivers list return an error
		consumerType := sdk.AuthConsumerType(vars["consumerType"])
		if !consumerType.IsValid() {
			return sdk.WithStack(sdk.ErrNotFound)
		}
		driver, ok := api.AuthenticationDrivers[consumerType]
		if !ok {
			return sdk.WithStack(sdk.ErrNotFound)
		}

		// Extract and validate signin request
		var req sdk.AuthConsumerSigninRequest
		if err := service.UnmarshalBody(r, &req); err != nil {
			return err
		}
		if err := driver.CheckSigninRequest(req); err != nil {
			return err
		}

		// Extract and validate signin state
		switch x := driver.(type) {
		case sdk.AuthDriverWithSigninStateToken:
			if err := x.CheckSigninStateToken(req); err != nil {
				return err
			}
		}

		// Convert code to external user info
		userInfo, err := driver.GetUserInfo(ctx, req)
		if err != nil {
			return err
		}

		tx, err := api.mustDB().Begin()
		if err != nil {
			return sdk.WithStack(err)
		}
		defer tx.Rollback() // nolint

		var signupDone bool
		initToken := req.String("init_token")
		hasInitToken := initToken != ""

		// Check if a consumer exists for consumer type and external user identifier
		consumer, err := authentication.LoadConsumerByTypeAndUserExternalID(ctx, tx, consumerType, userInfo.ExternalID,
			authentication.LoadConsumerOptions.WithAuthentifiedUser)
		if err != nil && !sdk.ErrorIs(err, sdk.ErrNotFound) {
			return err
		}

		// If there is no existing consumer we should check if a user need to be created
		// Then we want to create a new consumer for current type
		var u *sdk.AuthentifiedUser
		if consumer != nil {
			u = consumer.AuthentifiedUser
		} else {
			currentConsumer := getAPIConsumer(ctx)
			if currentConsumer != nil {
				// If no consumer already exists for given request, but there is a current session
				// We should create a new consumer for the current consumer's user
				u = currentConsumer.AuthentifiedUser

				// If new consumer email not already on exiting user add it
				existingContact, err := user.LoadContactByTypeAndValue(ctx, tx, sdk.UserContactTypeEmail, userInfo.Email)
				if err != nil && !sdk.ErrorIs(err, sdk.ErrNotFound) {
					return err
				}
				if existingContact == nil {
					// Insert a secondary contact for the existing user in database
					if err := user.InsertContact(ctx, tx, &sdk.UserContact{
						Primary:  false,
						Type:     sdk.UserContactTypeEmail,
						UserID:   u.ID,
						Value:    userInfo.Email,
						Verified: true,
					}); err != nil {
						return err
					}
				}
				if !api.Config.Auth.DisableAddUserInDefaultGroup {
					if err := group.CheckUserInDefaultGroup(ctx, tx, u.ID); err != nil {
						return err
					}
				}
			} else {
				// Check if a user already exists for external username
				u, err = user.LoadByUsername(ctx, tx, userInfo.Username, user.LoadOptions.WithContacts)
				if err != nil && !sdk.ErrorIs(err, sdk.ErrUserNotFound) {
					return err
				}
				if u != nil {
					// If the user exists with the same email address than in the userInfo,
					// we will create a new consumer and continue the signin
					// else we raise an error
					if u.GetEmail() != userInfo.Email {
						return sdk.NewErrorFrom(sdk.ErrForbidden, "a user already exists for username %s", userInfo.Username)
					}
				} else {
					// Check if a user already exists for external email
					contact, err := user.LoadContactByTypeAndValue(ctx, tx, sdk.UserContactTypeEmail, userInfo.Email)
					if err != nil && !sdk.ErrorIs(err, sdk.ErrNotFound) {
						return err
					}
					if contact != nil {
						// A user already exists with an other username but the same email address
						u, err = user.LoadByID(ctx, tx, contact.UserID, user.LoadOptions.WithContacts)
						if err != nil {
							return err
						}
					} else {
						// We can't find any user with the same email address
						// So we will do signup for a new user from the data got from the auth driver
						if driver.GetManifest().SignupDisabled {
							return sdk.WithStack(sdk.ErrSignupDisabled)
						}

						u = &sdk.AuthentifiedUser{
							Ring:     sdk.UserRingUser,
							Username: userInfo.Username,
							Fullname: userInfo.Fullname,
						}

						// If a magic token is given and there is no admin already registered, set new user as admin
						countAdmins, err := user.CountAdmin(tx)
						if err != nil {
							return err
						}
						if countAdmins == 0 && hasInitToken {
							u.Ring = sdk.UserRingAdmin
						} else {
							hasInitToken = false
						}

						// Insert the new user in database
						if err := user.Insert(ctx, tx, u); err != nil {
							return err
						}

						// Insert the primary contact for the new user in database
						if err := user.InsertContact(ctx, tx, &sdk.UserContact{
							Primary:  true,
							Type:     sdk.UserContactTypeEmail,
							UserID:   u.ID,
							Value:    userInfo.Email,
							Verified: true,
						}); err != nil {
							return err
						}

						if !api.Config.Auth.DisableAddUserInDefaultGroup {
							if err := group.CheckUserInDefaultGroup(ctx, tx, u.ID); err != nil {
								return err
							}
						}
						signupDone = true
					}
				}
			}

			// Create a new consumer for the new user
			consumer, err = authentication.NewConsumerExternal(ctx, tx, u.ID, consumerType, userInfo)
			if err != nil {
				return err
			}
		}

		if err := api.userSetOrganization(ctx, tx, u, userInfo.Organization); err != nil {
			return err
		}

		// If a new user has been created and a first admin has been create,
		// let's init the builtin consumers from the magix token
		if signupDone && hasInitToken {
			if err := initBuiltinConsumersFromStartupConfig(ctx, tx, consumer, initToken); err != nil {
				return err
			}
		}

		// Generate a new session for consumer
		sessionDuration := driver.GetSessionDuration()
		var session *sdk.AuthSession
		if userInfo.MFA {
			session, err = authentication.NewSessionWithMFA(ctx, tx, api.Cache, consumer, sessionDuration)
		} else {
			session, err = authentication.NewSession(ctx, tx, consumer, sessionDuration)
		}
		if err != nil {
			return err
		}

		// Store the last authentication date on the consumer
		now := time.Now()
		consumer.LastAuthentication = &now
		if err := authentication.UpdateConsumerLastAuthentication(ctx, tx, consumer); err != nil {
			return err
		}

		log.Debug(ctx, "postAuthSigninHandler> new session %s created for %.2f seconds: %+v", session.ID, sessionDuration.Seconds(), session)

		// Generate a jwt for current session
		jwt, err := authentication.NewSessionJWT(session, userInfo.ExternalTokenID)
		if err != nil {
			return err
		}

		usr, err := user.LoadByID(ctx, tx, consumer.AuthentifiedUserID)
		if err != nil {
			return err
		}

		if err := tx.Commit(); err != nil {
			return sdk.WithStack(err)
		}

		// Set a cookie with the jwt token
		api.SetCookie(w, service.JWTCookieName, jwt, session.ExpireAt, true)

		// Prepare http response
		resp := sdk.AuthConsumerSigninResponse{
			Token:  jwt,
			User:   usr,
			APIURL: api.Config.URL.API,
		}

		return service.WriteJSON(w, resp, http.StatusOK)
	}
}

func (api *API) postAuthSignoutHandler() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		session := getAuthSession(ctx)

		if err := authentication.DeleteSessionByID(api.mustDB(), session.ID); err != nil {
			return err
		}

		// Delete the jwt cookie value
		api.UnsetCookie(w, service.JWTCookieName, true)

		return service.WriteJSON(w, nil, http.StatusOK)
	}
}

func (api *API) postAuthDetachHandler() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		vars := mux.Vars(r)

		// Extract consumer type from request, is invalid or not in api drivers list return an error
		consumerType := sdk.AuthConsumerType(vars["consumerType"])
		if !consumerType.IsValidExternal() {
			return sdk.WithStack(sdk.ErrNotFound)
		}
		_, ok := api.AuthenticationDrivers[consumerType]
		if !ok {
			return sdk.WithStack(sdk.ErrNotFound)
		}

		currentConsumer := getAPIConsumer(ctx)

		tx, err := api.mustDB().Begin()
		if err != nil {
			return sdk.WithStack(err)
		}
		defer tx.Rollback() // nolint

		consumer, err := authentication.LoadConsumerByTypeAndUserID(ctx, tx, consumerType, currentConsumer.AuthentifiedUserID)
		if err != nil {
			return err
		}

		if err := authentication.DeleteConsumerByID(tx, consumer.ID); err != nil {
			return err
		}

		if err := tx.Commit(); err != nil {
			return sdk.WithStack(err)
		}

		// If we just removed the current consumer, clean http cookie.
		if consumer.ID == currentConsumer.ID {
			api.UnsetCookie(w, service.JWTCookieName, true)
		}

		return service.WriteJSON(w, nil, http.StatusOK)
	}
}

func (api *API) getAuthMe() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		m := getAuthDriverManifest(ctx)
		c := getAPIConsumer(ctx)
		s := getAuthSession(ctx)
		if m == nil || c == nil || s == nil {
			return sdk.WithStack(sdk.ErrUnauthorized)
		}

		// Clean user and consumer aggregated data
		u := *c.AuthentifiedUser
		u.Groups = nil
		c.AuthentifiedUser = nil

		return service.WriteJSON(w, sdk.AuthCurrentConsumerResponse{
			User:           u,
			Consumer:       *c,
			Session:        *s,
			DriverManifest: *m,
		}, http.StatusOK)
	}
}

func (api *API) getAuthSession() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		vars := mux.Vars(r)
		sessionID := vars["sessionID"]

		if !isCDN(ctx) {
			return sdk.WrapError(sdk.ErrForbidden, "only CDN can call this route")
		}

		s, err := authentication.LoadSessionByID(ctx, api.mustDB(), sessionID)
		if err != nil {
			return err
		}

		c, err := authentication.LoadConsumerByID(ctx, api.mustDB(), s.ConsumerID, authentication.LoadConsumerOptions.WithAuthentifiedUserWithContacts, authentication.LoadConsumerOptions.WithConsumerGroups)
		if err != nil {
			return err
		}

		return service.WriteJSON(w, sdk.AuthCurrentConsumerResponse{
			Consumer: *c,
			Session:  *s,
		}, http.StatusOK)
	}
}
<commit_msg>chore(api): log new session with mfa (#6151)<commit_after>package api

import (
	"context"
	"net/http"
	"time"

	"github.com/gorilla/mux"
	"github.com/rockbears/log"

	"github.com/ovh/cds/engine/api/authentication"
	"github.com/ovh/cds/engine/api/group"
	"github.com/ovh/cds/engine/api/user"
	"github.com/ovh/cds/engine/service"
	"github.com/ovh/cds/sdk"
)

func (api *API) getAuthDriversHandler() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		drivers := make(sdk.AuthDriverManifests, 0, len(api.AuthenticationDrivers))

		for _, d := range api.AuthenticationDrivers {
			drivers = append(drivers, d.GetManifest())
		}

		countAdmins, err := user.CountAdmin(api.mustDB())
		if err != nil {
			return err
		}

		var response = sdk.AuthDriverResponse{
			IsFirstConnection: countAdmins == 0,
			Drivers:           drivers,
		}

		return service.WriteJSON(w, response, http.StatusOK)
	}
}

func (api *API) getAuthScopesHandler() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		return service.WriteJSON(w, api.Router.scopeDetails, http.StatusOK)
	}
}

func (api *API) getAuthAskSigninHandler() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		vars := mux.Vars(r)

		// Extract consumer type from request, is invalid or not in api drivers list return an error
		consumerType := sdk.AuthConsumerType(vars["consumerType"])
		if !consumerType.IsValid() {
			return sdk.WithStack(sdk.ErrNotFound)
		}
		driver, ok := api.AuthenticationDrivers[consumerType]
		if !ok {
			return sdk.WithStack(sdk.ErrNotFound)
		}

		driverRedirect, ok := driver.(sdk.AuthDriverWithRedirect)
		if !ok {
			return nil
		}

		countAdmins, err := user.CountAdmin(api.mustDB())
		if err != nil {
			return err
		}

		var signinState = sdk.AuthSigninConsumerToken{
			RequireMFA:        QueryBool(r, "require_mfa"),
			RedirectURI:       QueryString(r, "redirect_uri"),
			IsFirstConnection: countAdmins == 0,
		}
		// Get the origin from request if set
		signinState.Origin = QueryString(r, "origin")

		if signinState.Origin != "" && !(signinState.Origin == "cdsctl" || signinState.Origin == "ui") {
			return sdk.NewErrorFrom(sdk.ErrWrongRequest, "invalid given origin value")
		}

		// Redirect to the right signin page depending on the consumer type
		redirect, err := driverRedirect.GetSigninURI(signinState)
		if err != nil {
			return err
		}

		return service.WriteJSON(w, redirect, http.StatusOK)
	}
}

func (api *API) postAuthSigninHandler() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		vars := mux.Vars(r)

		// Extract consumer type from request, is invalid or not in api drivers list return an error
		consumerType := sdk.AuthConsumerType(vars["consumerType"])
		if !consumerType.IsValid() {
			return sdk.WithStack(sdk.ErrNotFound)
		}
		driver, ok := api.AuthenticationDrivers[consumerType]
		if !ok {
			return sdk.WithStack(sdk.ErrNotFound)
		}

		// Extract and validate signin request
		var req sdk.AuthConsumerSigninRequest
		if err := service.UnmarshalBody(r, &req); err != nil {
			return err
		}
		if err := driver.CheckSigninRequest(req); err != nil {
			return err
		}

		// Extract and validate signin state
		switch x := driver.(type) {
		case sdk.AuthDriverWithSigninStateToken:
			if err := x.CheckSigninStateToken(req); err != nil {
				return err
			}
		}

		// Convert code to external user info
		userInfo, err := driver.GetUserInfo(ctx, req)
		if err != nil {
			return err
		}

		tx, err := api.mustDB().Begin()
		if err != nil {
			return sdk.WithStack(err)
		}
		defer tx.Rollback() // nolint

		var signupDone bool
		initToken := req.String("init_token")
		hasInitToken := initToken != ""

		// Check if a consumer exists for consumer type and external user identifier
		consumer, err := authentication.LoadConsumerByTypeAndUserExternalID(ctx, tx, consumerType, userInfo.ExternalID,
			authentication.LoadConsumerOptions.WithAuthentifiedUser)
		if err != nil && !sdk.ErrorIs(err, sdk.ErrNotFound) {
			return err
		}

		// If there is no existing consumer we should check if a user need to be created
		// Then we want to create a new consumer for current type
		var u *sdk.AuthentifiedUser
		if consumer != nil {
			u = consumer.AuthentifiedUser
		} else {
			currentConsumer := getAPIConsumer(ctx)
			if currentConsumer != nil {
				// If no consumer already exists for given request, but there is a current session
				// We should create a new consumer for the current consumer's user
				u = currentConsumer.AuthentifiedUser

				// If new consumer email not already on exiting user add it
				existingContact, err := user.LoadContactByTypeAndValue(ctx, tx, sdk.UserContactTypeEmail, userInfo.Email)
				if err != nil && !sdk.ErrorIs(err, sdk.ErrNotFound) {
					return err
				}
				if existingContact == nil {
					// Insert a secondary contact for the existing user in database
					if err := user.InsertContact(ctx, tx, &sdk.UserContact{
						Primary:  false,
						Type:     sdk.UserContactTypeEmail,
						UserID:   u.ID,
						Value:    userInfo.Email,
						Verified: true,
					}); err != nil {
						return err
					}
				}
				if !api.Config.Auth.DisableAddUserInDefaultGroup {
					if err := group.CheckUserInDefaultGroup(ctx, tx, u.ID); err != nil {
						return err
					}
				}
			} else {
				// Check if a user already exists for external username
				u, err = user.LoadByUsername(ctx, tx, userInfo.Username, user.LoadOptions.WithContacts)
				if err != nil && !sdk.ErrorIs(err, sdk.ErrUserNotFound) {
					return err
				}
				if u != nil {
					// If the user exists with the same email address than in the userInfo,
					// we will create a new consumer and continue the signin
					// else we raise an error
					if u.GetEmail() != userInfo.Email {
						return sdk.NewErrorFrom(sdk.ErrForbidden, "a user already exists for username %s", userInfo.Username)
					}
				} else {
					// Check if a user already exists for external email
					contact, err := user.LoadContactByTypeAndValue(ctx, tx, sdk.UserContactTypeEmail, userInfo.Email)
					if err != nil && !sdk.ErrorIs(err, sdk.ErrNotFound) {
						return err
					}
					if contact != nil {
						// A user already exists with an other username but the same email address
						u, err = user.LoadByID(ctx, tx, contact.UserID, user.LoadOptions.WithContacts)
						if err != nil {
							return err
						}
					} else {
						// We can't find any user with the same email address
						// So we will do signup for a new user from the data got from the auth driver
						if driver.GetManifest().SignupDisabled {
							return sdk.WithStack(sdk.ErrSignupDisabled)
						}

						u = &sdk.AuthentifiedUser{
							Ring:     sdk.UserRingUser,
							Username: userInfo.Username,
							Fullname: userInfo.Fullname,
						}

						// If a magic token is given and there is no admin already registered, set new user as admin
						countAdmins, err := user.CountAdmin(tx)
						if err != nil {
							return err
						}
						if countAdmins == 0 && hasInitToken {
							u.Ring = sdk.UserRingAdmin
						} else {
							hasInitToken = false
						}

						// Insert the new user in database
						if err := user.Insert(ctx, tx, u); err != nil {
							return err
						}

						// Insert the primary contact for the new user in database
						if err := user.InsertContact(ctx, tx, &sdk.UserContact{
							Primary:  true,
							Type:     sdk.UserContactTypeEmail,
							UserID:   u.ID,
							Value:    userInfo.Email,
							Verified: true,
						}); err != nil {
							return err
						}

						if !api.Config.Auth.DisableAddUserInDefaultGroup {
							if err := group.CheckUserInDefaultGroup(ctx, tx, u.ID); err != nil {
								return err
							}
						}
						signupDone = true
					}
				}
			}

			// Create a new consumer for the new user
			consumer, err = authentication.NewConsumerExternal(ctx, tx, u.ID, consumerType, userInfo)
			if err != nil {
				return err
			}
		}

		if err := api.userSetOrganization(ctx, tx, u, userInfo.Organization); err != nil {
			return err
		}

		// If a new user has been created and a first admin has been create,
		// let's init the builtin consumers from the magix token
		if signupDone && hasInitToken {
			if err := initBuiltinConsumersFromStartupConfig(ctx, tx, consumer, initToken); err != nil {
				return err
			}
		}

		// Generate a new session for consumer
		sessionDuration := driver.GetSessionDuration()
		var session *sdk.AuthSession
		if userInfo.MFA {
			trackSudo(ctx, w)
			session, err = authentication.NewSessionWithMFA(ctx, tx, api.Cache, consumer, sessionDuration)
			if err == nil {
				log.Info(ctx, "starting new session %s with MFA for consumer %s", session.ID, consumer.ID)
			}
		} else {
			session, err = authentication.NewSession(ctx, tx, consumer, sessionDuration)
		}
		if err != nil {
			return err
		}

		// Store the last authentication date on the consumer
		now := time.Now()
		consumer.LastAuthentication = &now
		if err := authentication.UpdateConsumerLastAuthentication(ctx, tx, consumer); err != nil {
			return err
		}

		log.Debug(ctx, "postAuthSigninHandler> new session %s created for %.2f seconds: %+v", session.ID, sessionDuration.Seconds(), session)

		// Generate a jwt for current session
		jwt, err := authentication.NewSessionJWT(session, userInfo.ExternalTokenID)
		if err != nil {
			return err
		}

		usr, err := user.LoadByID(ctx, tx, consumer.AuthentifiedUserID)
		if err != nil {
			return err
		}

		if err := tx.Commit(); err != nil {
			return sdk.WithStack(err)
		}

		// Set a cookie with the jwt token
		api.SetCookie(w, service.JWTCookieName, jwt, session.ExpireAt, true)

		// Prepare http response
		resp := sdk.AuthConsumerSigninResponse{
			Token:  jwt,
			User:   usr,
			APIURL: api.Config.URL.API,
		}

		return service.WriteJSON(w, resp, http.StatusOK)
	}
}

func (api *API) postAuthSignoutHandler() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		session := getAuthSession(ctx)

		if err := authentication.DeleteSessionByID(api.mustDB(), session.ID); err != nil {
			return err
		}

		// Delete the jwt cookie value
		api.UnsetCookie(w, service.JWTCookieName, true)

		return service.WriteJSON(w, nil, http.StatusOK)
	}
}

func (api *API) postAuthDetachHandler() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		vars := mux.Vars(r)

		// Extract consumer type from request, is invalid or not in api drivers list return an error
		consumerType := sdk.AuthConsumerType(vars["consumerType"])
		if !consumerType.IsValidExternal() {
			return sdk.WithStack(sdk.ErrNotFound)
		}
		_, ok := api.AuthenticationDrivers[consumerType]
		if !ok {
			return sdk.WithStack(sdk.ErrNotFound)
		}

		currentConsumer := getAPIConsumer(ctx)

		tx, err := api.mustDB().Begin()
		if err != nil {
			return sdk.WithStack(err)
		}
		defer tx.Rollback() // nolint

		consumer, err := authentication.LoadConsumerByTypeAndUserID(ctx, tx, consumerType, currentConsumer.AuthentifiedUserID)
		if err != nil {
			return err
		}

		if err := authentication.DeleteConsumerByID(tx, consumer.ID); err != nil {
			return err
		}

		if err := tx.Commit(); err != nil {
			return sdk.WithStack(err)
		}

		// If we just removed the current consumer, clean http cookie.
		if consumer.ID == currentConsumer.ID {
			api.UnsetCookie(w, service.JWTCookieName, true)
		}

		return service.WriteJSON(w, nil, http.StatusOK)
	}
}

func (api *API) getAuthMe() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		m := getAuthDriverManifest(ctx)
		c := getAPIConsumer(ctx)
		s := getAuthSession(ctx)
		if m == nil || c == nil || s == nil {
			return sdk.WithStack(sdk.ErrUnauthorized)
		}

		// Clean user and consumer aggregated data
		u := *c.AuthentifiedUser
		u.Groups = nil
		c.AuthentifiedUser = nil

		return service.WriteJSON(w, sdk.AuthCurrentConsumerResponse{
			User:           u,
			Consumer:       *c,
			Session:        *s,
			DriverManifest: *m,
		}, http.StatusOK)
	}
}

func (api *API) getAuthSession() service.Handler {
	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
		vars := mux.Vars(r)
		sessionID := vars["sessionID"]

		if !isCDN(ctx) {
			return sdk.WrapError(sdk.ErrForbidden, "only CDN can call this route")
		}

		s, err := authentication.LoadSessionByID(ctx, api.mustDB(), sessionID)
		if err != nil {
			return err
		}

		c, err := authentication.LoadConsumerByID(ctx, api.mustDB(), s.ConsumerID, authentication.LoadConsumerOptions.WithAuthentifiedUserWithContacts, authentication.LoadConsumerOptions.WithConsumerGroups)
		if err != nil {
			return err
		}

		return service.WriteJSON(w, sdk.AuthCurrentConsumerResponse{
			Consumer: *c,
			Session:  *s,
		}, http.StatusOK)
	}
}
<|endoftext|>


<commit_before>package coquelicot

import (
	"os"
)

// Attachment contain info about directory, base mime type and all files saved.
type Attachment struct {
	OriginalFile *OriginalFile
	Dir          *DirManager
	Versions     map[string]FileManager
}

// Function receive root directory, original file, convertion parameters.
// Return Attachment saved.
func Create(storage string, ofile *OriginalFile, converts map[string]string) (*Attachment, error) {
	dm, err := CreateDir(storage, ofile.BaseMime)
	if err != nil {
		return nil, err
	}

	attachment := &Attachment{
		OriginalFile: ofile,
		Dir:          dm,
		Versions:     make(map[string]FileManager),
	}

	if ofile.BaseMime == "image" {
		converts["thumbnail"] = "120x90"
	}
	for version, convert_opt := range converts {
		fm, err := attachment.CreateVersion(version, convert_opt)
		if err != nil {
			return nil, err
		}

		attachment.Versions[version] = fm
	}

	return attachment, os.Remove(attachment.OriginalFile.Filepath)
}

// Directly save single version and return FileManager.
func (attachment *Attachment) CreateVersion(version string, convert string) (FileManager, error) {
	fm := NewFileManager(attachment.Dir, attachment.OriginalFile.BaseMime, version)
	fm.SetFilename(attachment.OriginalFile)

	if err := fm.Convert(attachment.OriginalFile.Filepath, convert); err != nil {
		return nil, err
	}

	return fm, nil
}

func (attachment *Attachment) ToJson() map[string]interface{} {
	data := make(map[string]interface{})
	data["type"] = attachment.OriginalFile.BaseMime
	data["dir"] = attachment.Dir.Path
	data["name"] = attachment.OriginalFile.Filename
	versions := make(map[string]interface{})
	for version, fm := range attachment.Versions {
		versions[version] = fm.ToJson()
	}
	data["versions"] = versions

	return data
}
<commit_msg>Do not generate thumbnail unless -convert flag set<commit_after>package coquelicot

import (
	"os"
)

// Attachment contain info about directory, base mime type and all files saved.
type Attachment struct {
	OriginalFile *OriginalFile
	Dir          *DirManager
	Versions     map[string]FileManager
}

// Function receive root directory, original file, convertion parameters.
// Return Attachment saved.
func Create(storage string, ofile *OriginalFile, converts map[string]string) (*Attachment, error) {
	dm, err := CreateDir(storage, ofile.BaseMime)
	if err != nil {
		return nil, err
	}

	attachment := &Attachment{
		OriginalFile: ofile,
		Dir:          dm,
		Versions:     make(map[string]FileManager),
	}

	if ofile.BaseMime == "image" {
		converts["thumbnail"] = "120x90"
	}

	makeVersion := func(a *Attachment, version, convert string) error {
		fm, err := attachment.CreateVersion(version, convert)
		if err != nil {
			return err
		}
		attachment.Versions[version] = fm
		return nil
	}

	if err := makeVersion(attachment, "original", ""); err != nil {
		return nil, err
	}

	if makeThumbnail {
		if err := makeVersion(attachment, "thumbnail", converts["thumbnail"]); err != nil {
			return nil, err
		}
	}

	return attachment, os.Remove(attachment.OriginalFile.Filepath)
}

// Directly save single version and return FileManager.
func (attachment *Attachment) CreateVersion(version string, convert string) (FileManager, error) {
	fm := NewFileManager(attachment.Dir, attachment.OriginalFile.BaseMime, version)
	fm.SetFilename(attachment.OriginalFile)

	if err := fm.Convert(attachment.OriginalFile.Filepath, convert); err != nil {
		return nil, err
	}

	return fm, nil
}

func (attachment *Attachment) ToJson() map[string]interface{} {
	data := make(map[string]interface{})
	data["type"] = attachment.OriginalFile.BaseMime
	data["dir"] = attachment.Dir.Path
	data["name"] = attachment.OriginalFile.Filename
	versions := make(map[string]interface{})
	for version, fm := range attachment.Versions {
		versions[version] = fm.ToJson()
	}
	data["versions"] = versions

	return data
}
<|endoftext|>


<commit_before>package persistence

import (
	"testing"

	_ "github.com/go-sql-driver/mysql"
	"github.com/victorspringer/trapAdvisor/database"
	ta "github.com/victorspringer/trapAdvisor/touristattraction"
)

func TestTouristAttractionRepository(t *testing.T) {
	var err error
	database.DB, err = database.Init("test")
	if err != nil {
		t.Errorf("%v", err)
		t.FailNow()
	}
	defer func() {
		_, err = database.DB.Exec("DROP DATABASE IF EXISTS trapAdvisor_test")
		if err != nil {
			t.Errorf("%v", err)
			t.FailNow()
		}
		database.DB.Close()
	}()

	type args struct {
		t                   *ta.TouristAttraction
		touristAttractionID int
		tripID              int
		namePart            string
	}
	tt := struct {
		name         string
		r            *touristAttractionRepository
		args         args
		wantName     string
		wantNamePart string
		wantErr      bool
	}{
		"ValidTouristAttractionInsertion",
		&touristAttractionRepository{},
		args{
			&ta.TouristAttraction{
				Name:      "Test",
				Location:  "Brazil",
				VisitDate: "2007-02-02 02:02:20",
				Rating:    5.5,
				Pros:      "Test",
				Cons:      "Test",
				TripID:    1,
			},
			2,
			1,
			"Eif",
		},
		"Test",
		"Torre Eiffel",
		false,
	}
	t.Run(tt.name, func(t *testing.T) {
		if err = tt.r.Store(tt.args.t); (err != nil) != tt.wantErr {
			t.Errorf("touristAttractionRepository.Store() error = %v, wantErr %v", err, tt.wantErr)
			return
		}
		ta, err := tt.r.Find(tt.args.touristAttractionID)
		if (err != nil) != tt.wantErr {
			t.Errorf("touristAttractionRepository.Find() error = %v, wantErr %v", err, tt.wantErr)
			return
		}
		if ta.Name != tt.wantName {
			t.Errorf("invalid result")
			return
		}
		tas, err := tt.r.FindByTripID(tt.args.tripID)
		if (err != nil) != tt.wantErr {
			t.Errorf("touristAttractionRepository.FindByTripID() error = %v, wantErr %v", err, tt.wantErr)
			return
		}
		if !touristAttractionIDInSlice(tt.args.touristAttractionID, tas) {
			t.Errorf("invalid result")
			return
		}
		tasNamePart, err := tt.r.FindByNamePart(tt.args.namePart)
		if (err != nil) != tt.wantErr {
			t.Errorf("touristAttractionRepository.FindByNamePart() error = %v, wantErr %v", err, tt.wantErr)
			return
		}
		if tasNamePart[0][0] != tt.wantNamePart {
			t.Errorf("invalid result")
			return
		}
		tasMostVisited, err := tt.r.FindMostVisited()
		if (err != nil) != tt.wantErr {
			t.Errorf("touristAttractionRepository.FindMostVisited() error = %v, wantErr %v", err, tt.wantErr)
			return
		}
		if tasMostVisited[0][0] != tt.wantName {
			t.Errorf("invalid result")
			return
		}
		tasBestRated, err := tt.r.FindBestRated()
		if (err != nil) != tt.wantErr {
			t.Errorf("touristAttractionRepository.FindBestRated() error = %v, wantErr %v", err, tt.wantErr)
			return
		}
		if tasBestRated[1][0] != tt.wantName {
			t.Errorf("invalid result")
			return
		}
	})
}

func touristAttractionIDInSlice(id int, list []*ta.TouristAttraction) bool {
	for _, b := range list {
		if b.ID == id {
			return true
		}
	}
	return false
}
<commit_msg>travis-ci<commit_after>package persistence

import (
	"testing"

	"fmt"

	_ "github.com/go-sql-driver/mysql"
	"github.com/victorspringer/trapAdvisor/database"
	ta "github.com/victorspringer/trapAdvisor/touristattraction"
)

func TestTouristAttractionRepository(t *testing.T) {
	var err error
	database.DB, err = database.Init("test")
	if err != nil {
		t.Errorf("%v", err)
		t.FailNow()
	}
	defer func() {
		_, err = database.DB.Exec("DROP DATABASE IF EXISTS trapAdvisor_test")
		if err != nil {
			t.Errorf("%v", err)
			t.FailNow()
		}
		database.DB.Close()
	}()

	type args struct {
		t                   *ta.TouristAttraction
		touristAttractionID int
		tripID              int
		namePart            string
	}
	tt := struct {
		name         string
		r            *touristAttractionRepository
		args         args
		wantName     string
		wantNamePart string
		wantErr      bool
	}{
		"ValidTouristAttractionInsertion",
		&touristAttractionRepository{},
		args{
			&ta.TouristAttraction{
				Name:      "Test",
				Location:  "Brazil",
				VisitDate: "2007-02-02 02:02:20",
				Rating:    5.5,
				Pros:      "Test",
				Cons:      "Test",
				TripID:    1,
			},
			2,
			1,
			"Eif",
		},
		"Test",
		"Torre Eiffel",
		false,
	}
	t.Run(tt.name, func(t *testing.T) {
		if err = tt.r.Store(tt.args.t); (err != nil) != tt.wantErr {
			t.Errorf("touristAttractionRepository.Store() error = %v, wantErr %v", err, tt.wantErr)
			return
		}
		ta, err := tt.r.Find(tt.args.touristAttractionID)
		if (err != nil) != tt.wantErr {
			t.Errorf("touristAttractionRepository.Find() error = %v, wantErr %v", err, tt.wantErr)
			return
		}
		if ta.Name != tt.wantName {
			t.Errorf("invalid result")
			return
		}
		tas, err := tt.r.FindByTripID(tt.args.tripID)
		if (err != nil) != tt.wantErr {
			t.Errorf("touristAttractionRepository.FindByTripID() error = %v, wantErr %v", err, tt.wantErr)
			return
		}
		if !touristAttractionIDInSlice(tt.args.touristAttractionID, tas) {
			t.Errorf("invalid result")
			return
		}
		tasNamePart, err := tt.r.FindByNamePart(tt.args.namePart)
		if (err != nil) != tt.wantErr {
			t.Errorf("touristAttractionRepository.FindByNamePart() error = %v, wantErr %v", err, tt.wantErr)
			return
		}
		if tasNamePart[0][0] != tt.wantNamePart {
			t.Errorf("invalid result")
			return
		}
		tasMostVisited, err := tt.r.FindMostVisited()
		if (err != nil) != tt.wantErr {
			t.Errorf("touristAttractionRepository.FindMostVisited() error = %v, wantErr %v", err, tt.wantErr)
			return
		}
		fmt.Println(tasMostVisited[0][0], tasMostVisited[0][1], tasMostVisited[1][0], tasMostVisited[1][1])
		if tasMostVisited[0][0] != tt.wantName {
			t.Errorf("invalid result")
			return
		}
		tasBestRated, err := tt.r.FindBestRated()
		if (err != nil) != tt.wantErr {
			t.Errorf("touristAttractionRepository.FindBestRated() error = %v, wantErr %v", err, tt.wantErr)
			return
		}
		if tasBestRated[1][0] != tt.wantName {
			t.Errorf("invalid result")
			return
		}
	})
}

func touristAttractionIDInSlice(id int, list []*ta.TouristAttraction) bool {
	for _, b := range list {
		if b.ID == id {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>// +build linux

package cpu

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"unicode"
)

// Get cpu statistics
func Get() (*Stats, error) {
	// Reference: man 5 proc, Documentation/filesystems/proc.txt in Linux source code
	file, err := os.Open("/proc/stat")
	if err != nil {
		return nil, err
	}
	defer file.Close()
	return collectCPUStats(file)
}

// Stats represents cpu statistics for linux
type Stats struct {
	User, Nice, System, Idle, Iowait, Irq, Softirq, Steal, Guest, GuestNice uint64
	Total, CPUCount, StatCount                                              uint64
}

type cpuStat struct {
	name string
	ptr  *uint64
}

func collectCPUStats(out io.Reader) (*Stats, error) {
	scanner := bufio.NewScanner(out)
	var cpu Stats

	cpuStats := []cpuStat{
		{"user", &cpu.User},
		{"nice", &cpu.Nice},
		{"system", &cpu.System},
		{"idle", &cpu.Idle},
		{"iowait", &cpu.Iowait},
		{"irq", &cpu.Irq},
		{"softirq", &cpu.Softirq},
		{"steal", &cpu.Steal},
		{"guest", &cpu.Guest},
		{"guest_nice", &cpu.GuestNice},
	}

	if !scanner.Scan() {
		return nil, fmt.Errorf("failed to scan /proc/stat")
	}

	valStrs := strings.Fields(scanner.Text())[1:]
	cpu.StatCount = len(valStrs)
	for i, valStr := range valStrs {
		val, err := strconv.ParseUint(valStr, 10, 64)
		if err != nil {
			return nil, fmt.Errorf("failed to scan %s from /proc/stat", cpuStats[i].name)
		}
		*cpuStats[i].ptr = val
		cpu.Total += val
	}

	// Since cpustat[CPUTIME_USER] includes cpustat[CPUTIME_GUEST], subtract the duplicated values from total.
	// https://github.com/torvalds/linux/blob/4ec9f7a18/kernel/sched/cputime.c#L151-L158
	cpu.Total -= cpu.Guest
	// cpustat[CPUTIME_NICE] includes cpustat[CPUTIME_GUEST_NICE]
	cpu.Total -= cpu.GuestNice

	for scanner.Scan() {
		line := scanner.Text()
		if line[0] == 'c' && line[1] == 'p' && line[2] == 'u' && unicode.IsDigit(rune(line[3])) {
			cpu.CPUCount++
		}
	}
	if err := scanner.Err(); err != nil {
		return nil, fmt.Errorf("scan error for /proc/stat: %s", err)
	}

	return &cpu, nil
}
<commit_msg>change CPUCount and StatCount to int for linux<commit_after>// +build linux

package cpu

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"unicode"
)

// Get cpu statistics
func Get() (*Stats, error) {
	// Reference: man 5 proc, Documentation/filesystems/proc.txt in Linux source code
	file, err := os.Open("/proc/stat")
	if err != nil {
		return nil, err
	}
	defer file.Close()
	return collectCPUStats(file)
}

// Stats represents cpu statistics for linux
type Stats struct {
	User, Nice, System, Idle, Iowait, Irq, Softirq, Steal, Guest, GuestNice, Total uint64
	CPUCount, StatCount                                                            int
}

type cpuStat struct {
	name string
	ptr  *uint64
}

func collectCPUStats(out io.Reader) (*Stats, error) {
	scanner := bufio.NewScanner(out)
	var cpu Stats

	cpuStats := []cpuStat{
		{"user", &cpu.User},
		{"nice", &cpu.Nice},
		{"system", &cpu.System},
		{"idle", &cpu.Idle},
		{"iowait", &cpu.Iowait},
		{"irq", &cpu.Irq},
		{"softirq", &cpu.Softirq},
		{"steal", &cpu.Steal},
		{"guest", &cpu.Guest},
		{"guest_nice", &cpu.GuestNice},
	}

	if !scanner.Scan() {
		return nil, fmt.Errorf("failed to scan /proc/stat")
	}

	valStrs := strings.Fields(scanner.Text())[1:]
	cpu.StatCount = len(valStrs)
	for i, valStr := range valStrs {
		val, err := strconv.ParseUint(valStr, 10, 64)
		if err != nil {
			return nil, fmt.Errorf("failed to scan %s from /proc/stat", cpuStats[i].name)
		}
		*cpuStats[i].ptr = val
		cpu.Total += val
	}

	// Since cpustat[CPUTIME_USER] includes cpustat[CPUTIME_GUEST], subtract the duplicated values from total.
	// https://github.com/torvalds/linux/blob/4ec9f7a18/kernel/sched/cputime.c#L151-L158
	cpu.Total -= cpu.Guest
	// cpustat[CPUTIME_NICE] includes cpustat[CPUTIME_GUEST_NICE]
	cpu.Total -= cpu.GuestNice

	for scanner.Scan() {
		line := scanner.Text()
		if line[0] == 'c' && line[1] == 'p' && line[2] == 'u' && unicode.IsDigit(rune(line[3])) {
			cpu.CPUCount++
		}
	}
	if err := scanner.Err(); err != nil {
		return nil, fmt.Errorf("scan error for /proc/stat: %s", err)
	}

	return &cpu, nil
}
<|endoftext|>