<commit_before>/*
 * Minio Cloud Storage, (C) 2016 Minio, Inc.
 *
 * 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 cmd

import (
	"net/http"
	"strings"
	"time"
)

// Validates the preconditions for CopyObject, returns true if CopyObject operation should not proceed.
// Preconditions supported are:
//  x-amz-copy-source-if-modified-since
//  x-amz-copy-source-if-unmodified-since
//  x-amz-copy-source-if-match
//  x-amz-copy-source-if-none-match
func checkCopyObjectPreconditions(w http.ResponseWriter, r *http.Request, objInfo ObjectInfo) bool {
	// Return false for methods other than GET and HEAD.
	if r.Method != "PUT" {
		return false
	}
	// If the object doesn't have a modtime (IsZero), or the modtime
	// is obviously garbage (Unix time == 0), then ignore modtimes
	// and don't process the If-Modified-Since header.
	if objInfo.ModTime.IsZero() || objInfo.ModTime.Equal(time.Unix(0, 0)) {
		return false
	}

	// Headers to be set of object content is not going to be written to the client.
	writeHeaders := func() {
		// set common headers
		setCommonHeaders(w)

		// set object-related metadata headers
		w.Header().Set("Last-Modified", objInfo.ModTime.UTC().Format(http.TimeFormat))

		if objInfo.MD5Sum != "" {
			w.Header().Set("ETag", "\""+objInfo.MD5Sum+"\"")
		}
	}
	// x-amz-copy-source-if-modified-since: Return the object only if it has been modified
	// since the specified time otherwise return 412 (precondition failed).
	ifModifiedSinceHeader := r.Header.Get("x-amz-copy-source-if-modified-since")
	if ifModifiedSinceHeader != "" {
		if !ifModifiedSince(objInfo.ModTime, ifModifiedSinceHeader) {
			// If the object is not modified since the specified time.
			writeHeaders()
			writeErrorResponse(w, r, ErrPreconditionFailed, r.URL.Path)
			return true
		}
	}

	// x-amz-copy-source-if-unmodified-since : Return the object only if it has not been
	// modified since the specified time, otherwise return a 412 (precondition failed).
	ifUnmodifiedSinceHeader := r.Header.Get("x-amz-copy-source-if-unmodified-since")
	if ifUnmodifiedSinceHeader != "" {
		if ifModifiedSince(objInfo.ModTime, ifUnmodifiedSinceHeader) {
			// If the object is modified since the specified time.
			writeHeaders()
			writeErrorResponse(w, r, ErrPreconditionFailed, r.URL.Path)
			return true
		}
	}

	// x-amz-copy-source-if-match : Return the object only if its entity tag (ETag) is the
	// same as the one specified; otherwise return a 412 (precondition failed).
	ifMatchETagHeader := r.Header.Get("x-amz-copy-source-if-match")
	if ifMatchETagHeader != "" {
		if !isETagEqual(objInfo.MD5Sum, ifMatchETagHeader) {
			// If the object ETag does not match with the specified ETag.
			writeHeaders()
			writeErrorResponse(w, r, ErrPreconditionFailed, r.URL.Path)
			return true
		}
	}

	// If-None-Match : Return the object only if its entity tag (ETag) is different from the
	// one specified otherwise, return a 304 (not modified).
	ifNoneMatchETagHeader := r.Header.Get("x-amz-copy-source-if-none-match")
	if ifNoneMatchETagHeader != "" {
		if isETagEqual(objInfo.MD5Sum, ifNoneMatchETagHeader) {
			// If the object ETag matches with the specified ETag.
			writeHeaders()
			writeErrorResponse(w, r, ErrPreconditionFailed, r.URL.Path)
			return true
		}
	}
	// Object content should be written to http.ResponseWriter
	return false
}

// Validates the preconditions. Returns true if GET/HEAD operation should not proceed.
// Preconditions supported are:
//  If-Modified-Since
//  If-Unmodified-Since
//  If-Match
//  If-None-Match
func checkPreconditions(w http.ResponseWriter, r *http.Request, objInfo ObjectInfo) bool {
	// Return false for methods other than GET and HEAD.
	if r.Method != "GET" && r.Method != "HEAD" {
		return false
	}
	// If the object doesn't have a modtime (IsZero), or the modtime
	// is obviously garbage (Unix time == 0), then ignore modtimes
	// and don't process the If-Modified-Since header.
	if objInfo.ModTime.IsZero() || objInfo.ModTime.Equal(time.Unix(0, 0)) {
		return false
	}

	// Headers to be set of object content is not going to be written to the client.
	writeHeaders := func() {
		// set common headers
		setCommonHeaders(w)

		// set object-related metadata headers
		w.Header().Set("Last-Modified", objInfo.ModTime.UTC().Format(http.TimeFormat))

		if objInfo.MD5Sum != "" {
			w.Header().Set("ETag", "\""+objInfo.MD5Sum+"\"")
		}
	}
	// If-Modified-Since : Return the object only if it has been modified since the specified time,
	// otherwise return a 304 (not modified).
	ifModifiedSinceHeader := r.Header.Get("If-Modified-Since")
	if ifModifiedSinceHeader != "" {
		if !ifModifiedSince(objInfo.ModTime, ifModifiedSinceHeader) {
			// If the object is not modified since the specified time.
			writeHeaders()
			w.WriteHeader(http.StatusNotModified)
			return true
		}
	}

	// If-Unmodified-Since : Return the object only if it has not been modified since the specified
	// time, otherwise return a 412 (precondition failed).
	ifUnmodifiedSinceHeader := r.Header.Get("If-Unmodified-Since")
	if ifUnmodifiedSinceHeader != "" {
		if ifModifiedSince(objInfo.ModTime, ifUnmodifiedSinceHeader) {
			// If the object is modified since the specified time.
			writeHeaders()
			writeErrorResponse(w, r, ErrPreconditionFailed, r.URL.Path)
			return true
		}
	}

	// If-Match : Return the object only if its entity tag (ETag) is the same as the one specified;
	// otherwise return a 412 (precondition failed).
	ifMatchETagHeader := r.Header.Get("If-Match")
	if ifMatchETagHeader != "" {
		if !isETagEqual(objInfo.MD5Sum, ifMatchETagHeader) {
			// If the object ETag does not match with the specified ETag.
			writeHeaders()
			writeErrorResponse(w, r, ErrPreconditionFailed, r.URL.Path)
			return true
		}
	}

	// If-None-Match : Return the object only if its entity tag (ETag) is different from the
	// one specified otherwise, return a 304 (not modified).
	ifNoneMatchETagHeader := r.Header.Get("If-None-Match")
	if ifNoneMatchETagHeader != "" {
		if isETagEqual(objInfo.MD5Sum, ifNoneMatchETagHeader) {
			// If the object ETag matches with the specified ETag.
			writeHeaders()
			w.WriteHeader(http.StatusNotModified)
			return true
		}
	}
	// Object content should be written to http.ResponseWriter
	return false
}

// returns true if object was modified after givenTime.
func ifModifiedSince(objTime time.Time, givenTimeStr string) bool {
	givenTime, err := time.Parse(http.TimeFormat, givenTimeStr)
	if err != nil {
		return true
	}
	// The Date-Modified header truncates sub-second precision, so
	// use mtime < t+1s instead of mtime <= t to check for unmodified.
	if objTime.After(givenTime.Add(1 * time.Second)) {
		return true
	}
	return false
}

// canonicalizeETag returns ETag with leading and trailing double-quotes removed,
// if any present
func canonicalizeETag(etag string) string {
	canonicalETag := strings.TrimPrefix(etag, "\"")
	return strings.TrimSuffix(canonicalETag, "\"")
}

// isETagEqual return true if the canonical representations of two ETag strings
// are equal, false otherwise
func isETagEqual(left, right string) bool {
	return canonicalizeETag(left) == canonicalizeETag(right)
}
<commit_msg>Ignore copy conditions when ETag is not available (#2888)<commit_after>/*
 * Minio Cloud Storage, (C) 2016 Minio, Inc.
 *
 * 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 cmd

import (
	"net/http"
	"strings"
	"time"
)

// Validates the preconditions for CopyObject, returns true if CopyObject operation should not proceed.
// Preconditions supported are:
//  x-amz-copy-source-if-modified-since
//  x-amz-copy-source-if-unmodified-since
//  x-amz-copy-source-if-match
//  x-amz-copy-source-if-none-match
func checkCopyObjectPreconditions(w http.ResponseWriter, r *http.Request, objInfo ObjectInfo) bool {
	// Return false for methods other than GET and HEAD.
	if r.Method != "PUT" {
		return false
	}
	// If the object doesn't have a modtime (IsZero), or the modtime
	// is obviously garbage (Unix time == 0), then ignore modtimes
	// and don't process the If-Modified-Since header.
	if objInfo.ModTime.IsZero() || objInfo.ModTime.Equal(time.Unix(0, 0)) {
		return false
	}

	// Headers to be set of object content is not going to be written to the client.
	writeHeaders := func() {
		// set common headers
		setCommonHeaders(w)

		// set object-related metadata headers
		w.Header().Set("Last-Modified", objInfo.ModTime.UTC().Format(http.TimeFormat))

		if objInfo.MD5Sum != "" {
			w.Header().Set("ETag", "\""+objInfo.MD5Sum+"\"")
		}
	}
	// x-amz-copy-source-if-modified-since: Return the object only if it has been modified
	// since the specified time otherwise return 412 (precondition failed).
	ifModifiedSinceHeader := r.Header.Get("x-amz-copy-source-if-modified-since")
	if ifModifiedSinceHeader != "" {
		if !ifModifiedSince(objInfo.ModTime, ifModifiedSinceHeader) {
			// If the object is not modified since the specified time.
			writeHeaders()
			writeErrorResponse(w, r, ErrPreconditionFailed, r.URL.Path)
			return true
		}
	}

	// x-amz-copy-source-if-unmodified-since : Return the object only if it has not been
	// modified since the specified time, otherwise return a 412 (precondition failed).
	ifUnmodifiedSinceHeader := r.Header.Get("x-amz-copy-source-if-unmodified-since")
	if ifUnmodifiedSinceHeader != "" {
		if ifModifiedSince(objInfo.ModTime, ifUnmodifiedSinceHeader) {
			// If the object is modified since the specified time.
			writeHeaders()
			writeErrorResponse(w, r, ErrPreconditionFailed, r.URL.Path)
			return true
		}
	}

	// x-amz-copy-source-if-match : Return the object only if its entity tag (ETag) is the
	// same as the one specified; otherwise return a 412 (precondition failed).
	ifMatchETagHeader := r.Header.Get("x-amz-copy-source-if-match")
	if ifMatchETagHeader != "" {
		if objInfo.MD5Sum != "" && !isETagEqual(objInfo.MD5Sum, ifMatchETagHeader) {
			// If the object ETag does not match with the specified ETag.
			writeHeaders()
			writeErrorResponse(w, r, ErrPreconditionFailed, r.URL.Path)
			return true
		}
	}

	// If-None-Match : Return the object only if its entity tag (ETag) is different from the
	// one specified otherwise, return a 304 (not modified).
	ifNoneMatchETagHeader := r.Header.Get("x-amz-copy-source-if-none-match")
	if ifNoneMatchETagHeader != "" {
		if objInfo.MD5Sum != "" && isETagEqual(objInfo.MD5Sum, ifNoneMatchETagHeader) {
			// If the object ETag matches with the specified ETag.
			writeHeaders()
			writeErrorResponse(w, r, ErrPreconditionFailed, r.URL.Path)
			return true
		}
	}
	// Object content should be written to http.ResponseWriter
	return false
}

// Validates the preconditions. Returns true if GET/HEAD operation should not proceed.
// Preconditions supported are:
//  If-Modified-Since
//  If-Unmodified-Since
//  If-Match
//  If-None-Match
func checkPreconditions(w http.ResponseWriter, r *http.Request, objInfo ObjectInfo) bool {
	// Return false for methods other than GET and HEAD.
	if r.Method != "GET" && r.Method != "HEAD" {
		return false
	}
	// If the object doesn't have a modtime (IsZero), or the modtime
	// is obviously garbage (Unix time == 0), then ignore modtimes
	// and don't process the If-Modified-Since header.
	if objInfo.ModTime.IsZero() || objInfo.ModTime.Equal(time.Unix(0, 0)) {
		return false
	}

	// Headers to be set of object content is not going to be written to the client.
	writeHeaders := func() {
		// set common headers
		setCommonHeaders(w)

		// set object-related metadata headers
		w.Header().Set("Last-Modified", objInfo.ModTime.UTC().Format(http.TimeFormat))

		if objInfo.MD5Sum != "" {
			w.Header().Set("ETag", "\""+objInfo.MD5Sum+"\"")
		}
	}
	// If-Modified-Since : Return the object only if it has been modified since the specified time,
	// otherwise return a 304 (not modified).
	ifModifiedSinceHeader := r.Header.Get("If-Modified-Since")
	if ifModifiedSinceHeader != "" {
		if !ifModifiedSince(objInfo.ModTime, ifModifiedSinceHeader) {
			// If the object is not modified since the specified time.
			writeHeaders()
			w.WriteHeader(http.StatusNotModified)
			return true
		}
	}

	// If-Unmodified-Since : Return the object only if it has not been modified since the specified
	// time, otherwise return a 412 (precondition failed).
	ifUnmodifiedSinceHeader := r.Header.Get("If-Unmodified-Since")
	if ifUnmodifiedSinceHeader != "" {
		if ifModifiedSince(objInfo.ModTime, ifUnmodifiedSinceHeader) {
			// If the object is modified since the specified time.
			writeHeaders()
			writeErrorResponse(w, r, ErrPreconditionFailed, r.URL.Path)
			return true
		}
	}

	// If-Match : Return the object only if its entity tag (ETag) is the same as the one specified;
	// otherwise return a 412 (precondition failed).
	ifMatchETagHeader := r.Header.Get("If-Match")
	if ifMatchETagHeader != "" {
		if !isETagEqual(objInfo.MD5Sum, ifMatchETagHeader) {
			// If the object ETag does not match with the specified ETag.
			writeHeaders()
			writeErrorResponse(w, r, ErrPreconditionFailed, r.URL.Path)
			return true
		}
	}

	// If-None-Match : Return the object only if its entity tag (ETag) is different from the
	// one specified otherwise, return a 304 (not modified).
	ifNoneMatchETagHeader := r.Header.Get("If-None-Match")
	if ifNoneMatchETagHeader != "" {
		if isETagEqual(objInfo.MD5Sum, ifNoneMatchETagHeader) {
			// If the object ETag matches with the specified ETag.
			writeHeaders()
			w.WriteHeader(http.StatusNotModified)
			return true
		}
	}
	// Object content should be written to http.ResponseWriter
	return false
}

// returns true if object was modified after givenTime.
func ifModifiedSince(objTime time.Time, givenTimeStr string) bool {
	givenTime, err := time.Parse(http.TimeFormat, givenTimeStr)
	if err != nil {
		return true
	}
	// The Date-Modified header truncates sub-second precision, so
	// use mtime < t+1s instead of mtime <= t to check for unmodified.
	if objTime.After(givenTime.Add(1 * time.Second)) {
		return true
	}
	return false
}

// canonicalizeETag returns ETag with leading and trailing double-quotes removed,
// if any present
func canonicalizeETag(etag string) string {
	canonicalETag := strings.TrimPrefix(etag, "\"")
	return strings.TrimSuffix(canonicalETag, "\"")
}

// isETagEqual return true if the canonical representations of two ETag strings
// are equal, false otherwise
func isETagEqual(left, right string) bool {
	return canonicalizeETag(left) == canonicalizeETag(right)
}
<|endoftext|>


<commit_before>package mongoproto

import (
	"errors"
	"io"
)

var (
	ErrInvalidSize = errors.New("mongoproto: got invalid document size")
)

// CopyMessage copies reads & writes an entire message.
func CopyMessage(w io.Writer, r io.Reader) error {
	h, err := ReadHeader(r)
	if err != nil {
		return err
	}
	if err := h.WriteTo(w); err != nil {
		return err
	}
	_, err = io.CopyN(w, r, int64(h.MessageLength-MsgHeaderLen))
	return err
}

// ReadDocument read an entire BSON document. This document can be used with
// bson.Unmarshal.
func ReadDocument(r io.Reader) ([]byte, error) {
	var sizeRaw [4]byte
	if _, err := io.ReadFull(r, sizeRaw[:]); err != nil {
		return nil, err
	}
	size := getInt32(sizeRaw[:], 0)
	if size < 0 {
		return nil, ErrInvalidSize
	}
	doc := make([]byte, size)
	if size == 0 {
		return doc, nil
	}
	if size < 4 {
		return doc, nil
	}
	setInt32(doc, 0, size)

	if _, err := io.ReadFull(r, doc[4:]); err != nil {
		return nil, err
	}
	return doc, nil
}

// readCStringFromReader reads a null turminated string from an io.Reader.
func readCStringFromReader(r io.Reader) ([]byte, error) {
	var b []byte
	var n [1]byte
	for {
		if _, err := io.ReadFull(r, n[:]); err != nil {
			return nil, err
		}
		if n[0] == 0 {
			return b, nil
		}
		b = append(b, n[0])
	}
}

func readCString(b []byte) string {
	for i := 0; i < len(b); i++ {
		if b[i] == 0 {
			return string(b[:i])
		}
	}
	return ""
}

// all data in the MongoDB wire protocol is little-endian.
// all the read/write functions below are little-endian.
func getInt32(b []byte, pos int) int32 {
	return (int32(b[pos+0])) |
		(int32(b[pos+1]) << 8) |
		(int32(b[pos+2]) << 16) |
		(int32(b[pos+3]) << 24)
}

func setInt32(b []byte, pos int, i int32) {
	b[pos] = byte(i)
	b[pos+1] = byte(i >> 8)
	b[pos+2] = byte(i >> 16)
	b[pos+3] = byte(i >> 24)
}

func getInt64(b []byte, pos int) int64 {
	return (int64(b[pos+0])) |
		(int64(b[pos+1]) << 8) |
		(int64(b[pos+2]) << 16) |
		(int64(b[pos+3]) << 24) |
		(int64(b[pos+4]) << 32) |
		(int64(b[pos+5]) << 40) |
		(int64(b[pos+6]) << 48) |
		(int64(b[pos+7]) << 56)
}
<commit_msg>enforce 16mb limit<commit_after>package mongoproto

import (
	"errors"
	"io"
)

var (
	ErrInvalidSize = errors.New("mongoproto: got invalid document size")
)

const (
	maximumDocumentSize = 16 * 1024 * 1024 // 16MB max
)

// CopyMessage copies reads & writes an entire message.
func CopyMessage(w io.Writer, r io.Reader) error {
	h, err := ReadHeader(r)
	if err != nil {
		return err
	}
	if err := h.WriteTo(w); err != nil {
		return err
	}
	_, err = io.CopyN(w, r, int64(h.MessageLength-MsgHeaderLen))
	return err
}

// ReadDocument read an entire BSON document. This document can be used with
// bson.Unmarshal.
func ReadDocument(r io.Reader) ([]byte, error) {
	var sizeRaw [4]byte
	if _, err := io.ReadFull(r, sizeRaw[:]); err != nil {
		return nil, err
	}
	size := getInt32(sizeRaw[:], 0)
	if size < 0 {
		return nil, ErrInvalidSize
	}
	if size > maximumDocumentSize {
		return nil, ErrInvalidSize
	}
	doc := make([]byte, size)
	if size == 0 {
		return doc, nil
	}
	if size < 4 {
		return doc, nil
	}
	setInt32(doc, 0, size)

	if _, err := io.ReadFull(r, doc[4:]); err != nil {
		return nil, err
	}
	return doc, nil
}

// readCStringFromReader reads a null turminated string from an io.Reader.
func readCStringFromReader(r io.Reader) ([]byte, error) {
	var b []byte
	var n [1]byte
	for {
		if _, err := io.ReadFull(r, n[:]); err != nil {
			return nil, err
		}
		if n[0] == 0 {
			return b, nil
		}
		b = append(b, n[0])
	}
}

func readCString(b []byte) string {
	for i := 0; i < len(b); i++ {
		if b[i] == 0 {
			return string(b[:i])
		}
	}
	return ""
}

// all data in the MongoDB wire protocol is little-endian.
// all the read/write functions below are little-endian.
func getInt32(b []byte, pos int) int32 {
	return (int32(b[pos+0])) |
		(int32(b[pos+1]) << 8) |
		(int32(b[pos+2]) << 16) |
		(int32(b[pos+3]) << 24)
}

func setInt32(b []byte, pos int, i int32) {
	b[pos] = byte(i)
	b[pos+1] = byte(i >> 8)
	b[pos+2] = byte(i >> 16)
	b[pos+3] = byte(i >> 24)
}

func getInt64(b []byte, pos int) int64 {
	return (int64(b[pos+0])) |
		(int64(b[pos+1]) << 8) |
		(int64(b[pos+2]) << 16) |
		(int64(b[pos+3]) << 24) |
		(int64(b[pos+4]) << 32) |
		(int64(b[pos+5]) << 40) |
		(int64(b[pos+6]) << 48) |
		(int64(b[pos+7]) << 56)
}
<|endoftext|>


<commit_before>package mqtt

import (
	"bufio"
	"net"
)

// Connection represents a Network Connection.
type Connection struct {
	net.Conn
	R *bufio.Reader
	W *bufio.Writer
}

// NewConnection connects to the address on the named network,
// creates a Network Connection and returns it.
func NewConnection(network, address string) (*Connection, error) {
	// Connect to the address on the named network
	conn, err := net.Dial(network, address)
	if err != nil {
		return nil, err
	}

	// Create a Network Connection.
	c := &Connection{
		Conn: conn,
		R:    bufio.NewReader(conn),
		W:    bufio.NewWriter(conn),
	}

	return c, nil
}
<commit_msg>Update mqtt/connection.go<commit_after>package mqtt

import (
	"bufio"
	"net"
)

// Connection represents a Network Connection.
type Connection struct {
	net.Conn
	R *bufio.Reader
	W *bufio.Writer
}

// NewConnection connects to the address on the named network,
// creates a Network Connection and returns it.
func NewConnection(network, address string) (*Connection, error) {
	// Connect to the address on the named network.
	conn, err := net.Dial(network, address)
	if err != nil {
		return nil, err
	}

	// Create a Network Connection.
	c := &Connection{
		Conn: conn,
		R:    bufio.NewReader(conn),
		W:    bufio.NewWriter(conn),
	}

	// Return the Network Connection.
	return c, nil
}
<|endoftext|>


<commit_before>// Copyright 2018 Istio 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 workloadapi

import (
	"fmt"
	"log"

	rpc "github.com/gogo/googleapis/google/rpc"
	"golang.org/x/net/context"
	"google.golang.org/grpc"

	wlh "istio.io/istio/security/cmd/node_agent_k8s/workload/handler"
	pb "istio.io/istio/security/proto"
)

// server implements WorkloadService from workload_service.proto.
type server struct{}

// RegisterGrpc register grpc
func RegisterGrpc(s *grpc.Server) {
	pb.RegisterWorkloadServiceServer(s, &server{})
}

// Check do the check
func (s *server) Check(ctx context.Context, request *pb.CheckRequest) (*pb.CheckResponse, error) {

	log.Printf("[%v]: %v Check called", s, request)
	// Get the caller's credentials from the context.
	creds, e := wlh.CallerFromContext(ctx)
	if !e {
		resp := fmt.Sprint("Not able to get credentials")
		status := &rpc.Status{Code: int32(rpc.PERMISSION_DENIED), Message: resp}
		return &pb.CheckResponse{Status: status}, nil
	}

	log.Printf("Credentials are %v", creds)

	resp := fmt.Sprintf("all good to workload with service account %v", creds.ServiceAccount)
	status := &rpc.Status{Code: int32(rpc.OK), Message: resp}
	return &pb.CheckResponse{Status: status}, nil
}
<commit_msg>Fix lint. (#4453)<commit_after>// Copyright 2018 Istio 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 workloadapi

import (
	"fmt"
	"log"

	rpc "github.com/gogo/googleapis/google/rpc"
	"golang.org/x/net/context"
	"google.golang.org/grpc"

	wlh "istio.io/istio/security/cmd/node_agent_k8s/workload/handler"
	pb "istio.io/istio/security/proto"
)

// server implements WorkloadService from workload_service.proto.
type server struct{}

// RegisterGrpc register grpc
func RegisterGrpc(s *grpc.Server) {
	pb.RegisterWorkloadServiceServer(s, &server{})
}

// Check do the check
func (s *server) Check(ctx context.Context, request *pb.CheckRequest) (*pb.CheckResponse, error) {

	log.Printf("[%v]: %v Check called", s, request)
	// Get the caller's credentials from the context.
	creds, e := wlh.CallerFromContext(ctx)
	if !e {
		resp := "Not able to get credentials"
		status := &rpc.Status{Code: int32(rpc.PERMISSION_DENIED), Message: resp}
		return &pb.CheckResponse{Status: status}, nil
	}

	log.Printf("Credentials are %v", creds)

	resp := fmt.Sprintf("all good to workload with service account %v", creds.ServiceAccount)
	status := &rpc.Status{Code: int32(rpc.OK), Message: resp}
	return &pb.CheckResponse{Status: status}, nil
}
<|endoftext|>


<commit_before>package company

import (
	"net/http"

	"github.com/aerogo/aero"
	"github.com/animenotifier/arn"
	"github.com/animenotifier/notify.moe/components"
	"github.com/animenotifier/notify.moe/utils"
)

// Get company.
func Get(ctx *aero.Context) string {
	user := utils.GetUser(ctx)
	id := ctx.Get("id")
	company, err := arn.GetCompany(id)

	if err != nil {
		return ctx.Error(http.StatusNotFound, "Company not found", err)
	}

	openGraph := &arn.OpenGraph{
		Tags: map[string]string{
			"og:title":     company.Name.English,
			"og:url":       "https://" + ctx.App.Config.Domain + company.Link(),
			"og:site_name": "notify.moe",
			"og:type":      "article",
		},
	}

	if company.Image != "" {
		openGraph.Tags["og:image"] = company.Image
	}

	ctx.Data = openGraph
	return ctx.HTML(components.CompanyPage(company, user))
}
<commit_msg>Improved open graph data<commit_after>package company

import (
	"net/http"

	"github.com/aerogo/aero"
	"github.com/animenotifier/arn"
	"github.com/animenotifier/notify.moe/components"
	"github.com/animenotifier/notify.moe/utils"
)

// Get company.
func Get(ctx *aero.Context) string {
	user := utils.GetUser(ctx)
	id := ctx.Get("id")
	company, err := arn.GetCompany(id)

	if err != nil {
		return ctx.Error(http.StatusNotFound, "Company not found", err)
	}

	openGraph := &arn.OpenGraph{
		Tags: map[string]string{
			"og:title":       company.Name.English,
			"og:description": company.Description,
			"og:url":         "https://" + ctx.App.Config.Domain + company.Link(),
			"og:site_name":   "notify.moe",
			"og:type":        "article",
		},
	}

	if company.Image != "" {
		openGraph.Tags["og:image"] = company.Image
	}

	if company.Description != "" {
		openGraph.Tags["og:description"] = company.Description
	} else {
		openGraph.Tags["og:description"] = company.Name.English + " company information."
	}

	ctx.Data = openGraph
	return ctx.HTML(components.CompanyPage(company, user))
}
<|endoftext|>


<commit_before>package redis_test

import (
	"fmt"
	"math/rand"
	"net"
	"strings"

	"testing"
	"time"

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

	"gopkg.in/redis.v3"
)

type clusterScenario struct {
	ports     []string
	nodeIds   []string
	processes map[string]*redisProcess
	clients   map[string]*redis.Client
}

func (s *clusterScenario) primary() *redis.Client {
	return s.clients[s.ports[0]]
}

func (s *clusterScenario) masters() []*redis.Client {
	result := make([]*redis.Client, 3)
	for pos, port := range s.ports[:3] {
		result[pos] = s.clients[port]
	}
	return result
}

func (s *clusterScenario) slaves() []*redis.Client {
	result := make([]*redis.Client, 3)
	for pos, port := range s.ports[3:] {
		result[pos] = s.clients[port]
	}
	return result
}

func (s *clusterScenario) clusterClient(opt *redis.ClusterOptions) *redis.ClusterClient {
	addrs := make([]string, len(s.ports))
	for i, port := range s.ports {
		addrs[i] = net.JoinHostPort("127.0.0.1", port)
	}
	if opt == nil {
		opt = &redis.ClusterOptions{}
	}
	opt.Addrs = addrs
	return redis.NewClusterClient(opt)
}

func startCluster(scenario *clusterScenario) error {
	// Start processes and collect node ids
	for pos, port := range scenario.ports {
		process, err := startRedis(port, "--cluster-enabled", "yes")
		if err != nil {
			return err
		}

		client := redis.NewClient(&redis.Options{Addr: "127.0.0.1:" + port})
		info, err := client.ClusterNodes().Result()
		if err != nil {
			return err
		}

		scenario.processes[port] = process
		scenario.clients[port] = client
		scenario.nodeIds[pos] = info[:40]
	}

	// Meet cluster nodes
	for _, client := range scenario.clients {
		err := client.ClusterMeet("127.0.0.1", scenario.ports[0]).Err()
		if err != nil {
			return err
		}
	}

	// Bootstrap masters
	slots := []int{0, 5000, 10000, 16384}
	for pos, master := range scenario.masters() {
		err := master.ClusterAddSlotsRange(slots[pos], slots[pos+1]-1).Err()
		if err != nil {
			return err
		}
	}

	// Bootstrap slaves
	for idx, slave := range scenario.slaves() {
		masterId := scenario.nodeIds[idx]

		// Wait until master is available
		err := eventually(func() error {
			s := slave.ClusterNodes().Val()
			wanted := masterId
			if !strings.Contains(s, wanted) {
				return fmt.Errorf("%q does not contain %q", s, wanted)
			}
			return nil
		}, 10*time.Second)
		if err != nil {
			return err
		}

		err = slave.ClusterReplicate(masterId).Err()
		if err != nil {
			return err
		}
	}

	// Wait until all nodes have consistent info
	for _, client := range scenario.clients {
		err := eventually(func() error {
			for _, masterId := range scenario.nodeIds[:3] {
				s := client.ClusterNodes().Val()
				wanted := "slave " + masterId
				if !strings.Contains(s, wanted) {
					return fmt.Errorf("%q does not contain %q", s, wanted)
				}
			}
			return nil
		}, 10*time.Second)
		if err != nil {
			return err
		}
	}

	return nil
}

func stopCluster(scenario *clusterScenario) error {
	for _, client := range scenario.clients {
		if err := client.Close(); err != nil {
			return err
		}
	}
	for _, process := range scenario.processes {
		if err := process.Close(); err != nil {
			return err
		}
	}
	return nil
}

//------------------------------------------------------------------------------

var _ = Describe("Cluster", func() {
	Describe("HashSlot", func() {

		It("should calculate hash slots", func() {
			tests := []struct {
				key  string
				slot int
			}{
				{"123456789", 12739},
				{"{}foo", 9500},
				{"foo{}", 5542},
				{"foo{}{bar}", 8363},
				{"", 10503},
				{"", 5176},
				{string([]byte{83, 153, 134, 118, 229, 214, 244, 75, 140, 37, 215, 215}), 5463},
			}
			rand.Seed(100)

			for _, test := range tests {
				Expect(redis.HashSlot(test.key)).To(Equal(test.slot), "for %s", test.key)
			}
		})

		It("should extract keys from tags", func() {
			tests := []struct {
				one, two string
			}{
				{"foo{bar}", "bar"},
				{"{foo}bar", "foo"},
				{"{user1000}.following", "{user1000}.followers"},
				{"foo{{bar}}zap", "{bar"},
				{"foo{bar}{zap}", "bar"},
			}

			for _, test := range tests {
				Expect(redis.HashSlot(test.one)).To(Equal(redis.HashSlot(test.two)), "for %s <-> %s", test.one, test.two)
			}
		})

	})

	Describe("Commands", func() {

		It("should CLUSTER SLOTS", func() {
			res, err := cluster.primary().ClusterSlots().Result()
			Expect(err).NotTo(HaveOccurred())
			Expect(res).To(HaveLen(3))
			Expect(res).To(ConsistOf([]redis.ClusterSlotInfo{
				{0, 4999, []string{"127.0.0.1:8220", "127.0.0.1:8223"}},
				{5000, 9999, []string{"127.0.0.1:8221", "127.0.0.1:8224"}},
				{10000, 16383, []string{"127.0.0.1:8222", "127.0.0.1:8225"}},
			}))
		})

		It("should CLUSTER NODES", func() {
			res, err := cluster.primary().ClusterNodes().Result()
			Expect(err).NotTo(HaveOccurred())
			Expect(len(res)).To(BeNumerically(">", 400))
		})

		It("should CLUSTER INFO", func() {
			res, err := cluster.primary().ClusterInfo().Result()
			Expect(err).NotTo(HaveOccurred())
			Expect(res).To(ContainSubstring("cluster_known_nodes:6"))
		})

	})

	Describe("Client", func() {
		var client *redis.ClusterClient

		BeforeEach(func() {
			client = cluster.clusterClient(nil)
		})

		AfterEach(func() {
			for _, client := range cluster.masters() {
				Expect(client.FlushDb().Err()).NotTo(HaveOccurred())
			}
			Expect(client.Close()).NotTo(HaveOccurred())
		})

		It("should GET/SET/DEL", func() {
			val, err := client.Get("A").Result()
			Expect(err).To(Equal(redis.Nil))
			Expect(val).To(Equal(""))

			val, err = client.Set("A", "VALUE", 0).Result()
			Expect(err).NotTo(HaveOccurred())
			Expect(val).To(Equal("OK"))

			val, err = client.Get("A").Result()
			Expect(err).NotTo(HaveOccurred())
			Expect(val).To(Equal("VALUE"))

			cnt, err := client.Del("A").Result()
			Expect(err).NotTo(HaveOccurred())
			Expect(cnt).To(Equal(int64(1)))
		})

		It("should follow redirects", func() {
			Expect(client.Set("A", "VALUE", 0).Err()).NotTo(HaveOccurred())

			slot := redis.HashSlot("A")
			Expect(client.SwapSlot(slot)).To(Equal([]string{"127.0.0.1:8224", "127.0.0.1:8221"}))

			val, err := client.Get("A").Result()
			Expect(err).NotTo(HaveOccurred())
			Expect(val).To(Equal("VALUE"))

			Eventually(func() []string {
				return client.SlotAddrs(slot)
			}, "5s").Should(Equal([]string{"127.0.0.1:8221", "127.0.0.1:8224"}))
		})

		It("should perform multi-pipelines", func() {
			slot := redis.HashSlot("A")
			Expect(client.SwapSlot(slot)).To(Equal([]string{"127.0.0.1:8224", "127.0.0.1:8221"}))

			pipe := client.Pipeline()
			defer pipe.Close()

			keys := []string{"A", "B", "C", "D", "E", "F", "G"}
			for i, key := range keys {
				pipe.Set(key, key+"_value", 0)
				pipe.Expire(key, time.Duration(i+1)*time.Hour)
			}
			for _, key := range keys {
				pipe.Get(key)
				pipe.TTL(key)
			}

			cmds, err := pipe.Exec()
			Expect(err).NotTo(HaveOccurred())
			Expect(cmds).To(HaveLen(28))
			Expect(cmds[14].(*redis.StringCmd).Val()).To(Equal("A_value"))
			Expect(cmds[15].(*redis.DurationCmd).Val()).To(BeNumerically("~", 1*time.Hour, time.Second))
			Expect(cmds[20].(*redis.StringCmd).Val()).To(Equal("D_value"))
			Expect(cmds[21].(*redis.DurationCmd).Val()).To(BeNumerically("~", 4*time.Hour, time.Second))
			Expect(cmds[26].(*redis.StringCmd).Val()).To(Equal("G_value"))
			Expect(cmds[27].(*redis.DurationCmd).Val()).To(BeNumerically("~", 7*time.Hour, time.Second))
		})

		It("should return error when there are no attempts left", func() {
			Expect(client.Close()).NotTo(HaveOccurred())
			client = cluster.clusterClient(&redis.ClusterOptions{
				MaxRedirects: -1,
			})

			slot := redis.HashSlot("A")
			Expect(client.SwapSlot(slot)).To(Equal([]string{"127.0.0.1:8224", "127.0.0.1:8221"}))

			err := client.Get("A").Err()
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("MOVED"))
		})
	})
})

//------------------------------------------------------------------------------

func BenchmarkRedisClusterPing(b *testing.B) {
	if testing.Short() {
		b.Skip("skipping in short mode")
	}

	cluster := &clusterScenario{
		ports:     []string{"8220", "8221", "8222", "8223", "8224", "8225"},
		nodeIds:   make([]string, 6),
		processes: make(map[string]*redisProcess, 6),
		clients:   make(map[string]*redis.Client, 6),
	}
	if err := startCluster(cluster); err != nil {
		b.Fatal(err)
	}
	defer stopCluster(cluster)
	client := cluster.clusterClient(nil)
	defer client.Close()

	b.ResetTimer()

	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			if err := client.Ping().Err(); err != nil {
				b.Fatal(err)
			}
		}
	})
}
<commit_msg>tests: check Cluster node flags.<commit_after>package redis_test

import (
	"fmt"
	"math/rand"
	"net"
	"strings"

	"testing"
	"time"

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

	"gopkg.in/redis.v3"
)

type clusterScenario struct {
	ports     []string
	nodeIds   []string
	processes map[string]*redisProcess
	clients   map[string]*redis.Client
}

func (s *clusterScenario) primary() *redis.Client {
	return s.clients[s.ports[0]]
}

func (s *clusterScenario) masters() []*redis.Client {
	result := make([]*redis.Client, 3)
	for pos, port := range s.ports[:3] {
		result[pos] = s.clients[port]
	}
	return result
}

func (s *clusterScenario) slaves() []*redis.Client {
	result := make([]*redis.Client, 3)
	for pos, port := range s.ports[3:] {
		result[pos] = s.clients[port]
	}
	return result
}

func (s *clusterScenario) clusterClient(opt *redis.ClusterOptions) *redis.ClusterClient {
	addrs := make([]string, len(s.ports))
	for i, port := range s.ports {
		addrs[i] = net.JoinHostPort("127.0.0.1", port)
	}
	if opt == nil {
		opt = &redis.ClusterOptions{}
	}
	opt.Addrs = addrs
	return redis.NewClusterClient(opt)
}

func startCluster(scenario *clusterScenario) error {
	// Start processes and collect node ids
	for pos, port := range scenario.ports {
		process, err := startRedis(port, "--cluster-enabled", "yes")
		if err != nil {
			return err
		}

		client := redis.NewClient(&redis.Options{Addr: "127.0.0.1:" + port})
		info, err := client.ClusterNodes().Result()
		if err != nil {
			return err
		}

		scenario.processes[port] = process
		scenario.clients[port] = client
		scenario.nodeIds[pos] = info[:40]
	}

	// Meet cluster nodes
	for _, client := range scenario.clients {
		err := client.ClusterMeet("127.0.0.1", scenario.ports[0]).Err()
		if err != nil {
			return err
		}
	}

	// Bootstrap masters
	slots := []int{0, 5000, 10000, 16384}
	for pos, master := range scenario.masters() {
		err := master.ClusterAddSlotsRange(slots[pos], slots[pos+1]-1).Err()
		if err != nil {
			return err
		}
	}

	// Bootstrap slaves
	for idx, slave := range scenario.slaves() {
		masterId := scenario.nodeIds[idx]

		// Wait until master is available
		err := eventually(func() error {
			s := slave.ClusterNodes().Val()
			wanted := masterId
			if !strings.Contains(s, wanted) {
				return fmt.Errorf("%q does not contain %q", s, wanted)
			}
			return nil
		}, 10*time.Second)
		if err != nil {
			return err
		}

		err = slave.ClusterReplicate(masterId).Err()
		if err != nil {
			return err
		}
	}

	// Wait until all nodes have consistent info
	for _, client := range scenario.clients {
		err := eventually(func() error {
			s := client.ClusterNodes().Val()
			nodes := strings.Split(s, "\n")
			if len(nodes) < 6 {
				return fmt.Errorf("got %d nodes, wanted 6", len(nodes))
			}
			for _, node := range nodes {
				if node == "" {
					continue
				}
				parts := strings.Split(node, " ")
				var flags string
				if len(parts) >= 3 {
					flags = parts[2]
				}
				if !strings.Contains(flags, "master") && !strings.Contains(flags, "slave") {
					return fmt.Errorf("node flags are %q", flags)
				}
			}
			return nil
		}, 10*time.Second)
		if err != nil {
			return err
		}
	}

	return nil
}

func stopCluster(scenario *clusterScenario) error {
	for _, client := range scenario.clients {
		if err := client.Close(); err != nil {
			return err
		}
	}
	for _, process := range scenario.processes {
		if err := process.Close(); err != nil {
			return err
		}
	}
	return nil
}

//------------------------------------------------------------------------------

var _ = Describe("Cluster", func() {
	Describe("HashSlot", func() {

		It("should calculate hash slots", func() {
			tests := []struct {
				key  string
				slot int
			}{
				{"123456789", 12739},
				{"{}foo", 9500},
				{"foo{}", 5542},
				{"foo{}{bar}", 8363},
				{"", 10503},
				{"", 5176},
				{string([]byte{83, 153, 134, 118, 229, 214, 244, 75, 140, 37, 215, 215}), 5463},
			}
			rand.Seed(100)

			for _, test := range tests {
				Expect(redis.HashSlot(test.key)).To(Equal(test.slot), "for %s", test.key)
			}
		})

		It("should extract keys from tags", func() {
			tests := []struct {
				one, two string
			}{
				{"foo{bar}", "bar"},
				{"{foo}bar", "foo"},
				{"{user1000}.following", "{user1000}.followers"},
				{"foo{{bar}}zap", "{bar"},
				{"foo{bar}{zap}", "bar"},
			}

			for _, test := range tests {
				Expect(redis.HashSlot(test.one)).To(Equal(redis.HashSlot(test.two)), "for %s <-> %s", test.one, test.two)
			}
		})

	})

	Describe("Commands", func() {

		It("should CLUSTER SLOTS", func() {
			res, err := cluster.primary().ClusterSlots().Result()
			Expect(err).NotTo(HaveOccurred())
			Expect(res).To(HaveLen(3))
			Expect(res).To(ConsistOf([]redis.ClusterSlotInfo{
				{0, 4999, []string{"127.0.0.1:8220", "127.0.0.1:8223"}},
				{5000, 9999, []string{"127.0.0.1:8221", "127.0.0.1:8224"}},
				{10000, 16383, []string{"127.0.0.1:8222", "127.0.0.1:8225"}},
			}))
		})

		It("should CLUSTER NODES", func() {
			res, err := cluster.primary().ClusterNodes().Result()
			Expect(err).NotTo(HaveOccurred())
			Expect(len(res)).To(BeNumerically(">", 400))
		})

		It("should CLUSTER INFO", func() {
			res, err := cluster.primary().ClusterInfo().Result()
			Expect(err).NotTo(HaveOccurred())
			Expect(res).To(ContainSubstring("cluster_known_nodes:6"))
		})

	})

	Describe("Client", func() {
		var client *redis.ClusterClient

		BeforeEach(func() {
			client = cluster.clusterClient(nil)
		})

		AfterEach(func() {
			for _, client := range cluster.masters() {
				Expect(client.FlushDb().Err()).NotTo(HaveOccurred())
			}
			Expect(client.Close()).NotTo(HaveOccurred())
		})

		It("should GET/SET/DEL", func() {
			val, err := client.Get("A").Result()
			Expect(err).To(Equal(redis.Nil))
			Expect(val).To(Equal(""))

			val, err = client.Set("A", "VALUE", 0).Result()
			Expect(err).NotTo(HaveOccurred())
			Expect(val).To(Equal("OK"))

			val, err = client.Get("A").Result()
			Expect(err).NotTo(HaveOccurred())
			Expect(val).To(Equal("VALUE"))

			cnt, err := client.Del("A").Result()
			Expect(err).NotTo(HaveOccurred())
			Expect(cnt).To(Equal(int64(1)))
		})

		It("should follow redirects", func() {
			Expect(client.Set("A", "VALUE", 0).Err()).NotTo(HaveOccurred())

			slot := redis.HashSlot("A")
			Expect(client.SwapSlot(slot)).To(Equal([]string{"127.0.0.1:8224", "127.0.0.1:8221"}))

			val, err := client.Get("A").Result()
			Expect(err).NotTo(HaveOccurred())
			Expect(val).To(Equal("VALUE"))

			Eventually(func() []string {
				return client.SlotAddrs(slot)
			}, "5s").Should(Equal([]string{"127.0.0.1:8221", "127.0.0.1:8224"}))
		})

		It("should perform multi-pipelines", func() {
			slot := redis.HashSlot("A")
			Expect(client.SwapSlot(slot)).To(Equal([]string{"127.0.0.1:8224", "127.0.0.1:8221"}))

			pipe := client.Pipeline()
			defer pipe.Close()

			keys := []string{"A", "B", "C", "D", "E", "F", "G"}
			for i, key := range keys {
				pipe.Set(key, key+"_value", 0)
				pipe.Expire(key, time.Duration(i+1)*time.Hour)
			}
			for _, key := range keys {
				pipe.Get(key)
				pipe.TTL(key)
			}

			cmds, err := pipe.Exec()
			Expect(err).NotTo(HaveOccurred())
			Expect(cmds).To(HaveLen(28))
			Expect(cmds[14].(*redis.StringCmd).Val()).To(Equal("A_value"))
			Expect(cmds[15].(*redis.DurationCmd).Val()).To(BeNumerically("~", 1*time.Hour, time.Second))
			Expect(cmds[20].(*redis.StringCmd).Val()).To(Equal("D_value"))
			Expect(cmds[21].(*redis.DurationCmd).Val()).To(BeNumerically("~", 4*time.Hour, time.Second))
			Expect(cmds[26].(*redis.StringCmd).Val()).To(Equal("G_value"))
			Expect(cmds[27].(*redis.DurationCmd).Val()).To(BeNumerically("~", 7*time.Hour, time.Second))
		})

		It("should return error when there are no attempts left", func() {
			Expect(client.Close()).NotTo(HaveOccurred())
			client = cluster.clusterClient(&redis.ClusterOptions{
				MaxRedirects: -1,
			})

			slot := redis.HashSlot("A")
			Expect(client.SwapSlot(slot)).To(Equal([]string{"127.0.0.1:8224", "127.0.0.1:8221"}))

			err := client.Get("A").Err()
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("MOVED"))
		})
	})
})

//------------------------------------------------------------------------------

func BenchmarkRedisClusterPing(b *testing.B) {
	if testing.Short() {
		b.Skip("skipping in short mode")
	}

	cluster := &clusterScenario{
		ports:     []string{"8220", "8221", "8222", "8223", "8224", "8225"},
		nodeIds:   make([]string, 6),
		processes: make(map[string]*redisProcess, 6),
		clients:   make(map[string]*redis.Client, 6),
	}
	if err := startCluster(cluster); err != nil {
		b.Fatal(err)
	}
	defer stopCluster(cluster)
	client := cluster.clusterClient(nil)
	defer client.Close()

	b.ResetTimer()

	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			if err := client.Ping().Err(); err != nil {
				b.Fatal(err)
			}
		}
	})
}
<|endoftext|>


<commit_before>// Copyright 2015 Reborndb Org. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package main

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

	"github.com/reborndb/go/atomic2"
	"github.com/reborndb/reborn/pkg/models"
	"github.com/reborndb/reborn/pkg/utils"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/ngaut/zkhelper"
)

// HA mechanism
// 1, use cooridnator to elect a leader
// 2, leader check all servers (using ping command) every 3 seconds
// 3, if a server is down, leader will ask other agents to help it ping this server
// 4, if all agents reply this server is down, we will be sure it is down.
// 5, if the server is slave, we will only mark it offline.
// 6, if the server is master, we will do failover.
//
// if any step fails, we will discard current check and do above again.

type haTask struct {
	quited atomic2.Int64

	wg sync.WaitGroup
}

func (t *haTask) Run() error {
	log.Infof("we are leader now, begin check")

	t.quited.Set(0)
	t.wg.Add(1)
	defer t.wg.Done()

	for {
		if t.quited.Get() == 1 {
			break
		}

		if err := t.check(); err != nil {
			log.Errorf("check err %v", errors.ErrorStack(err))
			return errors.Trace(err)
		}

		// check servers every second
		time.Sleep(1 * time.Second)
	}

	return nil
}

func (t *haTask) Stop() {
	t.quited.Set(1)

	t.wg.Wait()

	log.Infof("stop ha check")
}

func (t *haTask) Interrupted() bool {
	return false
}

func (t *haTask) check() error {
	groups, err := models.ServerGroups(globalConn, globalEnv.ProductName())
	if err != nil {
		return errors.Trace(err)
	}

	cnt := 0

	const checkChanSize = 100
	ch := make(chan interface{}, checkChanSize)

	// check all servers in all groups
	for _, group := range groups {
		servers, err := group.GetServers(globalConn)
		if err != nil {
			return errors.Trace(err)
		}

		for _, server := range servers {
			cnt++

			go t.checkGroupServer(server, ch)
		}
	}

	var crashSlaves []*models.Server
	var crashMasters []*models.Server
	for i := 0; i < cnt; i++ {
		v := <-ch
		switch s := v.(type) {
		case nil:
		case *models.Server:
			if s.Type == models.SERVER_TYPE_SLAVE {
				crashSlaves = append(crashSlaves, s)
			} else if s.Type == models.SERVER_TYPE_MASTER {
				crashMasters = append(crashMasters, s)
			}
		case error:
			err = errors.Trace(s)
		}
	}

	if err != nil {
		return errors.Trace(err)
	}

	for _, s := range crashSlaves {
		log.Infof("slave %s in group %d is down, set offline", s.Addr, s.GroupId)
		group := models.NewServerGroup(globalEnv.ProductName(), s.GroupId)

		s.Type = models.SERVER_TYPE_OFFLINE

		if err := group.AddServer(globalConn, s, globalEnv.StoreAuth()); err != nil {
			return errors.Trace(err)
		}
	}

	for _, s := range crashMasters {
		log.Infof("master %s in group %d is down, do failover", s.Addr, s.GroupId)

		if err := t.doFailover(s); err != nil {
			log.Errorf("master %s in group %d is down, do failover err: %v", s.Addr, s.GroupId, err)
		}
	}

	return nil
}

func (t *haTask) checkGroupServer(s *models.Server, ch chan<- interface{}) {
	// we don't check offline server
	if s.Type == models.SERVER_TYPE_OFFLINE {
		ch <- nil
		return
	}

	var err error
	for i := 0; i < haMaxRetryNum; i++ {
		if err = utils.Ping(s.Addr, globalEnv.StoreAuth()); err == nil {
			break
		}

		err = errors.Trace(err)
		time.Sleep(time.Duration(haRetryDelay) * time.Second)
	}

	// here means we cannot ping server ok, so we think it is down
	if err == nil {
		ch <- nil
		return
	}

	log.Infof("leader check server %s in group %d err %v, let other agents help check", s.Addr, s.GroupId, err)

	// get all agents
	agents, err := getAgents()
	if err != nil {
		log.Errorf("get agents err %v", err)
		ch <- errors.Trace(err)
		return
	}

	reply := make([]interface{}, len(agents))

	var wg sync.WaitGroup
	for i, agent := range agents {
		if agent.ID == agentID {
			// ignore itself
			reply[i] = nil
			continue
		}

		wg.Add(1)
		go func(i int) {
			defer wg.Done()

			resp, err := http.Get(fmt.Sprintf("http://%s/api/check_store?addr=%s", agent.Addr, s.Addr))
			if err != nil {
				reply[i] = errors.Trace(err)
				return
			}
			defer resp.Body.Close()
			if _, err = ioutil.ReadAll(resp.Body); err != nil {
				reply[i] = errors.Trace(err)
				return
			}

			reply[i] = int(resp.StatusCode)
		}(i)
	}

	wg.Wait()

	for i, r := range reply {
		switch v := r.(type) {
		case nil:
			// itself, ignore
		case error:
			log.Errorf("let agent %s check %s err %v", agents[i].Addr, s.Addr, v)
			ch <- errors.Trace(v)
			return
		case int:
			if v == http.StatusOK {
				log.Infof("agent %s check %s ok, maybe it is alive", agents[i].Addr, s.Addr)
				ch <- nil
				return
			}
			// here mean agent check server failed
		}
	}

	// if all nodes check the store server is down, we will think it is down
	log.Infof("all agents check server %s is down", s.Addr)
	ch <- s
}

func (t *haTask) doFailover(s *models.Server) error {
	// first get all slaves
	group := models.NewServerGroup(globalEnv.ProductName(), s.GroupId)

	var err error
	group.Servers, err = group.GetServers(globalConn)
	if err != nil {
		return errors.Trace(err)
	}

	slaves := make([]*models.Server, 0, len(group.Servers))
	slaveAddrs := make([]string, 0, len(group.Servers))

	for _, s := range group.Servers {
		if s.Type == models.SERVER_TYPE_SLAVE {
			slaves = append(slaves, s)
			slaveAddrs = append(slaveAddrs, s.Addr)
		}
	}

	// elect a new master
	log.Infof("elect a new master in %v", slaveAddrs)
	addr, err := t.electNewMaster(slaves)
	if err != nil {
		return errors.Trace(err)
	}

	// prmote it as new master
	log.Infof("promote %s as the new master", addr)
	if err := group.Promote(globalConn, addr, globalEnv.StoreAuth()); err != nil {
		// should we fatal here and let human intervention ???
		return errors.Trace(err)
	}

	// let other slaves replicate from the new master
	for _, slave := range slaves {
		if slave.Addr == addr {
			continue
		}

		log.Infof("let %s slaveof new master %s", slave.Addr, addr)
		if err := utils.SlaveOf(slave.Addr, addr, globalEnv.StoreAuth()); err != nil {
			// should we fatal here and let human intervention ???
			return errors.Trace(err)
		}
	}

	return nil
}

func (t *haTask) getReplicationInfo(s *models.Server) (map[string]string, error) {
	v, err := utils.GetRedisInfo(s.Addr, "REPLICATION", globalEnv.StoreAuth())
	if err != nil {
		return nil, errors.Trace(err)
	}

	seps := strings.Split(v, "\r\n")
	// skip first line, is # Replication
	seps = seps[1:]

	m := make(map[string]string, len(seps))
	for _, s := range seps {
		kv := strings.SplitN(s, ":", 2)
		if len(kv) == 2 {
			m[strings.TrimSpace(kv[0])] = strings.TrimSpace(kv[1])
		}
	}
	return m, nil
}

func (t *haTask) electNewMaster(slaves []*models.Server) (string, error) {
	var addr string
	var checkOffset int64 = 0
	var checkPriority int = 0

	for _, slave := range slaves {
		m, err := t.getReplicationInfo(slave)
		if err != nil {
			return "", errors.Errorf("slave %s get replication info err %v", slave.Addr, err)
		}

		if m["slave"] == "master" {
			return "", errors.Errorf("server %s is not slave now,", slave.Addr)
		}

		if m["master_link_status"] == "up" {
			return "", errors.Errorf("slave %s master_link_status is up, master may be alive", slave.Addr)
		}

		priority, _ := strconv.Atoi(m["slave_priority"])
		replOffset, _ := strconv.ParseInt(m["slave_repl_offset"], 10, 64)

		used := false
		// like redis-sentinel, first check priority, then salve repl offset
		if checkPriority < priority {
			used = true
		} else if checkPriority == priority {
			if checkOffset < replOffset {
				used = true
			}
		}

		if used {
			addr = slave.Addr
			checkPriority = priority
			checkOffset = replOffset
		}
	}

	if len(addr) == 0 {
		return "", errors.Errorf("no proper candidate to be promoted")
	}

	log.Infof("select slave %s as new master, priority:%d, repl_offset:%d", addr, checkPriority, checkOffset)

	return addr, nil
}

func startHA() {
	elector := zkhelper.CreateElection(globalConn, fmt.Sprintf("/zk/reborn/db_%s/ha", globalEnv.ProductName()))
	task := &haTask{}

	err := elector.RunTask(task)
	if err != nil {
		log.Errorf("run elector task err: %v", err)
	}
}
<commit_msg>check err after all other agents checking returns<commit_after>// Copyright 2015 Reborndb Org. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package main

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

	"github.com/reborndb/go/atomic2"
	"github.com/reborndb/reborn/pkg/models"
	"github.com/reborndb/reborn/pkg/utils"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/ngaut/zkhelper"
)

// HA mechanism
// 1, use cooridnator to elect a leader
// 2, leader check all servers (using ping command) every 3 seconds
// 3, if a server is down, leader will ask other agents to help it ping this server
// 4, if all agents reply this server is down, we will be sure it is down.
// 5, if the server is slave, we will only mark it offline.
// 6, if the server is master, we will do failover.
//
// if any step fails, we will discard current check and do above again.

type haTask struct {
	quited atomic2.Int64

	wg sync.WaitGroup
}

func (t *haTask) Run() error {
	log.Infof("we are leader now, begin check")

	t.quited.Set(0)
	t.wg.Add(1)
	defer t.wg.Done()

	for {
		if t.quited.Get() == 1 {
			break
		}

		if err := t.check(); err != nil {
			log.Errorf("check err %v", errors.ErrorStack(err))
			return errors.Trace(err)
		}

		// check servers every second
		time.Sleep(1 * time.Second)
	}

	return nil
}

func (t *haTask) Stop() {
	t.quited.Set(1)

	t.wg.Wait()

	log.Infof("stop ha check")
}

func (t *haTask) Interrupted() bool {
	return false
}

func (t *haTask) check() error {
	groups, err := models.ServerGroups(globalConn, globalEnv.ProductName())
	if err != nil {
		return errors.Trace(err)
	}

	cnt := 0

	const checkChanSize = 100
	ch := make(chan interface{}, checkChanSize)

	// check all servers in all groups
	for _, group := range groups {
		servers, err := group.GetServers(globalConn)
		if err != nil {
			return errors.Trace(err)
		}

		for _, server := range servers {
			cnt++

			go t.checkGroupServer(server, ch)
		}
	}

	var crashSlaves []*models.Server
	var crashMasters []*models.Server
	for i := 0; i < cnt; i++ {
		v := <-ch
		switch s := v.(type) {
		case nil:
		case *models.Server:
			if s.Type == models.SERVER_TYPE_SLAVE {
				crashSlaves = append(crashSlaves, s)
			} else if s.Type == models.SERVER_TYPE_MASTER {
				crashMasters = append(crashMasters, s)
			}
		case error:
			err = errors.Trace(s)
		}
	}

	if err != nil {
		return errors.Trace(err)
	}

	for _, s := range crashSlaves {
		log.Infof("slave %s in group %d is down, set offline", s.Addr, s.GroupId)
		group := models.NewServerGroup(globalEnv.ProductName(), s.GroupId)

		s.Type = models.SERVER_TYPE_OFFLINE

		if err := group.AddServer(globalConn, s, globalEnv.StoreAuth()); err != nil {
			return errors.Trace(err)
		}
	}

	for _, s := range crashMasters {
		log.Infof("master %s in group %d is down, do failover", s.Addr, s.GroupId)

		if err := t.doFailover(s); err != nil {
			log.Errorf("master %s in group %d is down, do failover err: %v", s.Addr, s.GroupId, err)
		}
	}

	return nil
}

func (t *haTask) checkGroupServer(s *models.Server, ch chan<- interface{}) {
	// we don't check offline server
	if s.Type == models.SERVER_TYPE_OFFLINE {
		ch <- nil
		return
	}

	var err error
	for i := 0; i < haMaxRetryNum; i++ {
		if err = utils.Ping(s.Addr, globalEnv.StoreAuth()); err == nil {
			break
		}

		err = errors.Trace(err)
		time.Sleep(time.Duration(haRetryDelay) * time.Second)
	}

	if err == nil {
		ch <- nil
		return
	}

	// here means we cannot ping server ok, so we think it is down
	// let other help use to check
	log.Infof("leader check server %s in group %d err %v, let other agents help check", s.Addr, s.GroupId, err)

	// get all agents
	agents, err := getAgents()
	if err != nil {
		log.Errorf("get agents err %v", err)
		ch <- errors.Trace(err)
		return
	}

	reply := make([]interface{}, len(agents))

	var wg sync.WaitGroup
	for i, agent := range agents {
		if agent.ID == agentID {
			// ignore itself
			reply[i] = nil
			continue
		}

		wg.Add(1)
		go func(i int) {
			defer wg.Done()

			resp, err := http.Get(fmt.Sprintf("http://%s/api/check_store?addr=%s", agent.Addr, s.Addr))
			if err != nil {
				reply[i] = errors.Trace(err)
				return
			}
			defer resp.Body.Close()
			if _, err = ioutil.ReadAll(resp.Body); err != nil {
				reply[i] = errors.Trace(err)
				return
			}

			reply[i] = int(resp.StatusCode)
		}(i)
	}

	wg.Wait()

	for i, r := range reply {
		switch v := r.(type) {
		case nil:
			// itself, ignore
		case error:
			log.Errorf("let agent %s check %s err %v", agents[i].Addr, s.Addr, v)
			err = errors.Trace(v)
		case int:
			if v == http.StatusOK {
				log.Infof("agent %s check %s ok, maybe it is alive", agents[i].Addr, s.Addr)
				ch <- nil
				return
			}
			// here mean agent check server failed
		}
	}

	if err != nil {
		// here mean let some agent check err, maybe we cann't connect the agent
		// so return error to retry again
		ch <- errors.Trace(err)
		return
	}

	// if all nodes check the store server is down, we will think it is down
	log.Infof("all agents check server %s is down", s.Addr)
	ch <- s
}

func (t *haTask) doFailover(s *models.Server) error {
	// first get all slaves
	group := models.NewServerGroup(globalEnv.ProductName(), s.GroupId)

	var err error
	group.Servers, err = group.GetServers(globalConn)
	if err != nil {
		return errors.Trace(err)
	}

	slaves := make([]*models.Server, 0, len(group.Servers))
	slaveAddrs := make([]string, 0, len(group.Servers))

	for _, s := range group.Servers {
		if s.Type == models.SERVER_TYPE_SLAVE {
			slaves = append(slaves, s)
			slaveAddrs = append(slaveAddrs, s.Addr)
		}
	}

	// elect a new master
	log.Infof("elect a new master in %v", slaveAddrs)
	addr, err := t.electNewMaster(slaves)
	if err != nil {
		return errors.Trace(err)
	}

	// prmote it as new master
	log.Infof("promote %s as the new master", addr)
	if err := group.Promote(globalConn, addr, globalEnv.StoreAuth()); err != nil {
		// should we fatal here and let human intervention ???
		return errors.Trace(err)
	}

	// let other slaves replicate from the new master
	for _, slave := range slaves {
		if slave.Addr == addr {
			continue
		}

		log.Infof("let %s slaveof new master %s", slave.Addr, addr)
		if err := utils.SlaveOf(slave.Addr, addr, globalEnv.StoreAuth()); err != nil {
			// should we fatal here and let human intervention ???
			return errors.Trace(err)
		}
	}

	return nil
}

func (t *haTask) getReplicationInfo(s *models.Server) (map[string]string, error) {
	v, err := utils.GetRedisInfo(s.Addr, "REPLICATION", globalEnv.StoreAuth())
	if err != nil {
		return nil, errors.Trace(err)
	}

	seps := strings.Split(v, "\r\n")
	// skip first line, is # Replication
	seps = seps[1:]

	m := make(map[string]string, len(seps))
	for _, s := range seps {
		kv := strings.SplitN(s, ":", 2)
		if len(kv) == 2 {
			m[strings.TrimSpace(kv[0])] = strings.TrimSpace(kv[1])
		}
	}
	return m, nil
}

func (t *haTask) electNewMaster(slaves []*models.Server) (string, error) {
	var addr string
	var checkOffset int64 = 0
	var checkPriority int = 0

	for _, slave := range slaves {
		m, err := t.getReplicationInfo(slave)
		if err != nil {
			return "", errors.Errorf("slave %s get replication info err %v", slave.Addr, err)
		}

		if m["slave"] == "master" {
			return "", errors.Errorf("server %s is not slave now,", slave.Addr)
		}

		if m["master_link_status"] == "up" {
			return "", errors.Errorf("slave %s master_link_status is up, master may be alive", slave.Addr)
		}

		priority, _ := strconv.Atoi(m["slave_priority"])
		replOffset, _ := strconv.ParseInt(m["slave_repl_offset"], 10, 64)

		used := false
		// like redis-sentinel, first check priority, then salve repl offset
		if checkPriority < priority {
			used = true
		} else if checkPriority == priority {
			if checkOffset < replOffset {
				used = true
			}
		}

		if used {
			addr = slave.Addr
			checkPriority = priority
			checkOffset = replOffset
		}
	}

	if len(addr) == 0 {
		return "", errors.Errorf("no proper candidate to be promoted")
	}

	log.Infof("select slave %s as new master, priority:%d, repl_offset:%d", addr, checkPriority, checkOffset)

	return addr, nil
}

func startHA() {
	elector := zkhelper.CreateElection(globalConn, fmt.Sprintf("/zk/reborn/db_%s/ha", globalEnv.ProductName()))
	task := &haTask{}

	err := elector.RunTask(task)
	if err != nil {
		log.Errorf("run elector task err: %v", err)
	}
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2016 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 cmd

import (
	"fmt"
	"os"
	"os/signal"
	"syscall"

	"github.com/kardianos/osext"
	"github.com/spf13/cobra"

	"github.com/skydive-project/skydive/config"
	"github.com/skydive-project/skydive/logging"
)

var (
	analyzerUID uint32
)

// AllInOne skydive allinone root command
var AllInOne = &cobra.Command{
	Use:          "allinone",
	Short:        "Skydive All-In-One",
	Long:         "Skydive All-In-One (Analyzer + Agent)",
	SilenceUsage: true,
	Run: func(cmd *cobra.Command, args []string) {
		skydivePath, _ := osext.Executable()

		analyzerAttr := &os.ProcAttr{
			Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
			Env:   os.Environ(),
		}

		if analyzerUID != 0 {
			analyzerAttr.Sys = &syscall.SysProcAttr{
				Credential: &syscall.Credential{
					Uid: analyzerUID,
				},
			}
		}

		args = []string{"skydive"}

		for _, cfgFile := range CfgFiles {
			args = append(args, "-c")
			args = append(args, cfgFile)
		}

		if cfgBackend != "" {
			args = append(args, "-b")
			args = append(args, cfgBackend)
		}

		analyzerArgs := make([]string, len(args))
		copy(analyzerArgs, args)

		if len(CfgFiles) != 0 {
			if err := config.InitConfig(cfgBackend, CfgFiles); err != nil {
				panic(fmt.Sprintf("Failed to initialize config: %s", err.Error()))
			}
		}

		if err := logging.InitLogging(); err != nil {
			panic(fmt.Sprintf("Failed to initialize logging system: %s", err.Error()))
		}

		analyzer, err := os.StartProcess(skydivePath, append(analyzerArgs, "analyzer"), analyzerAttr)
		if err != nil {
			logging.GetLogger().Fatalf("Can't start Skydive All-in-One : %v", err)
		}

		os.Setenv("SKYDIVE_ANALYZERS", config.GetConfig().GetString("analyzer.listen"))

		agentAttr := &os.ProcAttr{
			Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
			Env:   os.Environ(),
		}

		agentArgs := make([]string, len(args))
		copy(agentArgs, args)

		agent, err := os.StartProcess(skydivePath, append(agentArgs, "agent"), agentAttr)
		if err != nil {
			logging.GetLogger().Fatalf("Can't start Skydive All-in-One : %v", err)
		}

		logging.GetLogger().Notice("Skydive All-in-One starting !")
		ch := make(chan os.Signal)
		signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
		<-ch

		analyzer.Kill()
		agent.Kill()

		analyzer.Wait()
		agent.Wait()

		logging.GetLogger().Notice("Skydive All-in-One stopped.")
	},
}

func init() {
	var uid uint32

	skydivePath, _ := osext.Executable()

	fi, err := os.Stat(skydivePath)
	if err == nil {
		uid = fi.Sys().(*syscall.Stat_t).Uid
	}

	AllInOne.Flags().Uint32VarP(&analyzerUID, "analyzer-uid", "", uid, "uid used for analyzer process")
}
<commit_msg>allinone: fix permission issue with file logging<commit_after>/*
 * Copyright (C) 2016 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 cmd

import (
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"

	"github.com/kardianos/osext"
	"github.com/spf13/cobra"

	"github.com/skydive-project/skydive/config"
	"github.com/skydive-project/skydive/logging"
)

var (
	analyzerUID uint32
)

// AllInOne skydive allinone root command
var AllInOne = &cobra.Command{
	Use:          "allinone",
	Short:        "Skydive All-In-One",
	Long:         "Skydive All-In-One (Analyzer + Agent)",
	SilenceUsage: true,
	Run: func(cmd *cobra.Command, args []string) {
		if uid := os.Geteuid(); uid != 0 {
			fmt.Fprintln(os.Stderr, "All-In-One mode has to be run as root")
			os.Exit(1)
		}

		skydivePath, _ := osext.Executable()
		logFile := config.GetConfig().GetString("logging.file.path")
		extension := filepath.Ext(logFile)
		logFile = strings.TrimSuffix(logFile, extension)
		os.Setenv("SKYDIVE_LOGGING_FILE_PATH", logFile+"-analyzer"+extension)

		analyzerAttr := &os.ProcAttr{
			Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
			Env:   os.Environ(),
		}

		if analyzerUID != 0 {
			analyzerAttr.Sys = &syscall.SysProcAttr{
				Credential: &syscall.Credential{
					Uid: analyzerUID,
				},
			}
		}

		args = []string{"skydive"}

		for _, cfgFile := range CfgFiles {
			args = append(args, "-c")
			args = append(args, cfgFile)
		}

		if cfgBackend != "" {
			args = append(args, "-b")
			args = append(args, cfgBackend)
		}

		analyzerArgs := make([]string, len(args))
		copy(analyzerArgs, args)

		if len(CfgFiles) != 0 {
			if err := config.InitConfig(cfgBackend, CfgFiles); err != nil {
				fmt.Fprintf(os.Stderr, "Failed to initialize config: %s", err.Error())
				os.Exit(1)
			}
		}

		analyzer, err := os.StartProcess(skydivePath, append(analyzerArgs, "analyzer"), analyzerAttr)
		if err != nil {
			logging.GetLogger().Errorf("Can't start Skydive analyzer: %v", err)
			os.Exit(1)
		}

		os.Setenv("SKYDIVE_ANALYZERS", config.GetConfig().GetString("analyzer.listen"))
		os.Setenv("SKYDIVE_LOGGING_FILE_PATH", logFile+"-agent"+extension)

		agentAttr := &os.ProcAttr{
			Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
			Env:   os.Environ(),
		}

		agentArgs := make([]string, len(args))
		copy(agentArgs, args)

		agent, err := os.StartProcess(skydivePath, append(agentArgs, "agent"), agentAttr)
		if err != nil {
			logging.GetLogger().Errorf("Can't start Skydive agent: %s. Please check you are root", err.Error())
			os.Exit(1)
		}

		logging.GetLogger().Notice("Skydive All-in-One starting !")
		ch := make(chan os.Signal)
		signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
		<-ch

		analyzer.Kill()
		agent.Kill()

		analyzer.Wait()
		agent.Wait()

		logging.GetLogger().Notice("Skydive All-in-One stopped.")
	},
}

func init() {
	var uid uint32

	skydivePath, _ := osext.Executable()

	fi, err := os.Stat(skydivePath)
	if err == nil {
		uid = fi.Sys().(*syscall.Stat_t).Uid
	}

	AllInOne.Flags().Uint32VarP(&analyzerUID, "analyzer-uid", "", uid, "uid used for analyzer process")
}
<|endoftext|>


<commit_before>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"

	"github.com/OWASP/Amass/v3/config"
	"github.com/OWASP/Amass/v3/format"
	"github.com/OWASP/Amass/v3/graph"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/OWASP/Amass/v3/stringset"
	"github.com/fatih/color"
)

const (
	dbUsageMsg = "db [options]"
)

type dbArgs struct {
	Domains stringset.Set
	Enum    int
	Options struct {
		DemoMode         bool
		IPs              bool
		IPv4             bool
		IPv6             bool
		ListEnumerations bool
		ASNTableSummary  bool
		DiscoveredNames  bool
		NoColor          bool
		ShowAll          bool
		Silent           bool
		Sources          bool
	}
	Filepaths struct {
		ConfigFile string
		Directory  string
		Domains    string
		JSONOutput string
		TermOut    string
	}
}

func runDBCommand(clArgs []string) {
	var args dbArgs
	var help1, help2 bool
	dbCommand := flag.NewFlagSet("db", flag.ContinueOnError)

	dbBuf := new(bytes.Buffer)
	dbCommand.SetOutput(dbBuf)

	args.Domains = stringset.New()

	dbCommand.BoolVar(&help1, "h", false, "Show the program usage message")
	dbCommand.BoolVar(&help2, "help", false, "Show the program usage message")
	dbCommand.Var(&args.Domains, "d", "Domain names separated by commas (can be used multiple times)")
	dbCommand.IntVar(&args.Enum, "enum", 0, "Identify an enumeration via an index from the listing")
	dbCommand.BoolVar(&args.Options.DemoMode, "demo", false, "Censor output to make it suitable for demonstrations")
	dbCommand.BoolVar(&args.Options.IPs, "ip", false, "Show the IP addresses for discovered names")
	dbCommand.BoolVar(&args.Options.IPv4, "ipv4", false, "Show the IPv4 addresses for discovered names")
	dbCommand.BoolVar(&args.Options.IPv6, "ipv6", false, "Show the IPv6 addresses for discovered names")
	dbCommand.BoolVar(&args.Options.ListEnumerations, "list", false, "Numbered list of enums filtered on provided domains")
	dbCommand.BoolVar(&args.Options.Sources, "src", false, "Print data sources for the discovered names")
	dbCommand.BoolVar(&args.Options.ASNTableSummary, "summary", false, "Print Just ASN Table Summary")
	dbCommand.BoolVar(&args.Options.DiscoveredNames, "names", false, "Print Just Discovered Names")
	dbCommand.BoolVar(&args.Options.NoColor, "nocolor", false, "Disable colorized output")
	dbCommand.BoolVar(&args.Options.ShowAll, "show", false, "Print the results for the enumeration index + domains provided")
	dbCommand.BoolVar(&args.Options.Silent, "silent", false, "Disable all output during execution")
	dbCommand.StringVar(&args.Filepaths.ConfigFile, "config", "", "Path to the INI configuration file. Additional details below")
	dbCommand.StringVar(&args.Filepaths.Directory, "dir", "", "Path to the directory containing the graph database")
	dbCommand.StringVar(&args.Filepaths.Domains, "df", "", "Path to a file providing root domain names")
	dbCommand.StringVar(&args.Filepaths.JSONOutput, "json", "", "Path to the JSON output file")
	dbCommand.StringVar(&args.Filepaths.TermOut, "o", "", "Path to the text file containing terminal stdout/stderr")

	if len(clArgs) < 1 {
		commandUsage(dbUsageMsg, dbCommand, dbBuf)
		return
	}

	if err := dbCommand.Parse(clArgs); err != nil {
		r.Fprintf(color.Error, "%v\n", err)
		os.Exit(1)
	}
	if help1 || help2 {
		commandUsage(dbUsageMsg, dbCommand, dbBuf)
		return
	}

	if args.Options.NoColor {
		color.NoColor = true
	}
	if args.Options.Silent {
		color.Output = ioutil.Discard
		color.Error = ioutil.Discard
	}

	if args.Filepaths.Domains != "" {
		list, err := config.GetListFromFile(args.Filepaths.Domains)
		if err != nil {
			r.Fprintf(color.Error, "Failed to parse the domain names file: %v\n", err)
			return
		}
		args.Domains.InsertMany(list...)
	}

	cfg := new(config.Config)
	cfg.LocalDatabase = true
	// Check if a configuration file was provided, and if so, load the settings
	if err := config.AcquireConfig(args.Filepaths.Directory, args.Filepaths.ConfigFile, cfg); err == nil {
		if args.Filepaths.Directory == "" {
			args.Filepaths.Directory = cfg.Dir
		}
		if len(args.Domains) == 0 {
			args.Domains.InsertMany(cfg.Domains()...)
		}
	} else if args.Filepaths.ConfigFile != "" {
		r.Fprintf(color.Error, "Failed to load the configuration file: %v\n", err)
		os.Exit(1)
	}

	db := openGraphDatabase(args.Filepaths.Directory, cfg)
	if db == nil {
		r.Fprintln(color.Error, "Failed to connect with the database")
		os.Exit(1)
	}
	defer db.Close()

	// Create the in-memory graph database for events that have information in scope
	memDB, err := memGraphForScope(args.Domains.Slice(), db)
	if err != nil {
		r.Fprintln(color.Error, err.Error())
		os.Exit(1)
	}

	// Get all the UUIDs for events that have information in scope
	uuids := memDB.EventList()
	if len(uuids) == 0 {
		r.Fprintln(color.Error, "Failed to find the domains of interest in the database")
		os.Exit(1)
	}

	if args.Options.ListEnumerations {
		listEvents(uuids, memDB)
		return
	}

	if args.Options.ShowAll || args.Filepaths.JSONOutput != "" {
		args.Options.DiscoveredNames = true
		args.Options.ASNTableSummary = true
	}

	if !args.Options.DiscoveredNames && !args.Options.ASNTableSummary {
		commandUsage(dbUsageMsg, dbCommand, dbBuf)
		return
	}

	// Put the events in chronological order
	uuids, _, _ = orderedEvents(uuids, memDB)
	if len(uuids) == 0 {
		r.Fprintln(color.Error, "Failed to sort the events")
		os.Exit(1)
	}

	// Select the enumeration that the user specified
	if args.Enum > 0 && len(uuids) >= args.Enum {
		idx := len(uuids) - args.Enum

		uuids = []string{uuids[idx]}
	}

	if args.Options.ASNTableSummary {
		fgY.Fprintln(color.Error, "Could take a moment while acquiring AS network information")
		// Migrate the changes back to the persistent db
		if healASInfo(uuids, memDB) {
			memDB.MigrateEvents(db, uuids...)
		}
	}

	showEventData(&args, uuids, memDB)
}

func listEvents(uuids []string, db *graph.Graph) {
	events, earliest, latest := orderedEvents(uuids, db)
	// Check if the user has requested the list of enumerations
	for pos, idx := 0, len(events)-1; idx >= 0; idx-- {
		if pos != 0 {
			g.Println()
		}

		g.Printf("%d) %s -> %s: ", pos+1, earliest[idx].Format(timeFormat), latest[idx].Format(timeFormat))
		// Print out the scope for this enumeration
		for x, domain := range db.EventDomains(events[idx]) {
			if x != 0 {
				g.Print(", ")
			}
			g.Print(domain)
		}
		g.Println()
		pos++
	}
}

func showEventData(args *dbArgs, uuids []string, db *graph.Graph) {
	var total int
	var err error
	var outfile *os.File
	var discovered []*requests.Output
	domains := args.Domains.Slice()

	if args.Filepaths.TermOut != "" {
		outfile, err = os.OpenFile(args.Filepaths.TermOut, os.O_WRONLY|os.O_CREATE, 0644)
		if err != nil {
			r.Fprintf(color.Error, "Failed to open the text output file: %v\n", err)
			os.Exit(1)
		}
		defer func() {
			outfile.Sync()
			outfile.Close()
		}()
		outfile.Truncate(0)
		outfile.Seek(0, 0)
	}

	tags := make(map[string]int)
	asns := make(map[int]*format.ASNSummaryData)
	for _, out := range getEventOutput(uuids, db) {
		if len(domains) > 0 && !domainNameInScope(out.Name, domains) {
			continue
		}

		out.Addresses = format.DesiredAddrTypes(out.Addresses, args.Options.IPv4, args.Options.IPv6)
		if l := len(out.Addresses); (args.Options.IPs || args.Options.IPv4 || args.Options.IPv6) && l == 0 {
			continue
		} else if l > 0 {
			total++
			format.UpdateSummaryData(out, tags, asns)
		}

		source, name, ips := format.OutputLineParts(out, args.Options.Sources,
			args.Options.IPs || args.Options.IPv4 || args.Options.IPv6, args.Options.DemoMode)
		if ips != "" {
			ips = " " + ips
		}

		if args.Options.DiscoveredNames {
			var written bool
			if outfile != nil {
				fmt.Fprintf(outfile, "%s%s%s\n", source, name, ips)
				written = true
			}
			if args.Filepaths.JSONOutput != "" {
				discovered = append(discovered, out)
				written = true
			}
			if !written {
				fmt.Fprintf(color.Output, "%s%s%s\n", blue(source), green(name), yellow(ips))
			}
		}
	}

	if total == 0 {
		r.Println("No names were discovered")
		return
	}
	if args.Filepaths.JSONOutput != "" {
		writeJSON(args, uuids, discovered, db)
	}
	if args.Options.ASNTableSummary {
		var out io.Writer
		status := color.NoColor

		if outfile != nil {
			out = outfile
			color.NoColor = true
		} else if args.Options.ShowAll {
			out = color.Error
		} else {
			out = color.Output
		}

		format.FprintEnumerationSummary(out, total, tags, asns, args.Options.DemoMode)
		color.NoColor = status
	}
}

type jsonEvent struct {
	UUID   string `json:"uuid"`
	Start  string `json:"start"`
	Finish string `json:"finish"`
}

type jsonDomain struct {
	Domain string `json:"domain"`
	Total  int    `json:"total"`
	Names  []*requests.Output
}

type jsonOutput struct {
	Events  []*jsonEvent  `json:"events"`
	Domains []*jsonDomain `json:"domains"`
}

func writeJSON(args *dbArgs, uuids []string, assets []*requests.Output, db *graph.Graph) {
	var output jsonOutput

	// Add the event data to the JSON
	events, earliest, latest := orderedEvents(uuids, db)
	for i, uuid := range events {
		output.Events = append(output.Events, &jsonEvent{
			UUID:   uuid,
			Start:  earliest[i].Format(timeFormat),
			Finish: latest[i].Format(timeFormat),
		})
	}
	// Add the asset specific data
	for _, asset := range assets {
		var found bool
		var d *jsonDomain

		for _, domain := range output.Domains {
			if domain.Domain == asset.Domain {
				found = true
				d = domain
				break
			}
		}
		if !found {
			d = &jsonDomain{Domain: asset.Domain}

			output.Domains = append(output.Domains, d)
		}

		d.Total++
		d.Names = append(d.Names, asset)
	}

	jsonptr, err := os.OpenFile(args.Filepaths.JSONOutput, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		r.Fprintf(color.Error, "Failed to open the JSON output file: %v\n", err)
		return
	}
	// Remove previously stored data and encode the JSON
	jsonptr.Truncate(0)
	jsonptr.Seek(0, 0)
	json.NewEncoder(jsonptr).Encode(output)
	jsonptr.Sync()
	jsonptr.Close()
}
<commit_msg>fixed the proper JSON output option<commit_after>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"

	"github.com/OWASP/Amass/v3/config"
	"github.com/OWASP/Amass/v3/format"
	"github.com/OWASP/Amass/v3/graph"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/OWASP/Amass/v3/stringset"
	"github.com/fatih/color"
)

const (
	dbUsageMsg = "db [options]"
)

type dbArgs struct {
	Domains stringset.Set
	Enum    int
	Options struct {
		DemoMode         bool
		IPs              bool
		IPv4             bool
		IPv6             bool
		ListEnumerations bool
		ASNTableSummary  bool
		DiscoveredNames  bool
		NoColor          bool
		ShowAll          bool
		Silent           bool
		Sources          bool
	}
	Filepaths struct {
		ConfigFile string
		Directory  string
		Domains    string
		JSONOutput string
		TermOut    string
	}
}

func runDBCommand(clArgs []string) {
	var args dbArgs
	var help1, help2 bool
	dbCommand := flag.NewFlagSet("db", flag.ContinueOnError)

	dbBuf := new(bytes.Buffer)
	dbCommand.SetOutput(dbBuf)

	args.Domains = stringset.New()

	dbCommand.BoolVar(&help1, "h", false, "Show the program usage message")
	dbCommand.BoolVar(&help2, "help", false, "Show the program usage message")
	dbCommand.Var(&args.Domains, "d", "Domain names separated by commas (can be used multiple times)")
	dbCommand.IntVar(&args.Enum, "enum", 0, "Identify an enumeration via an index from the listing")
	dbCommand.BoolVar(&args.Options.DemoMode, "demo", false, "Censor output to make it suitable for demonstrations")
	dbCommand.BoolVar(&args.Options.IPs, "ip", false, "Show the IP addresses for discovered names")
	dbCommand.BoolVar(&args.Options.IPv4, "ipv4", false, "Show the IPv4 addresses for discovered names")
	dbCommand.BoolVar(&args.Options.IPv6, "ipv6", false, "Show the IPv6 addresses for discovered names")
	dbCommand.BoolVar(&args.Options.ListEnumerations, "list", false, "Numbered list of enums filtered on provided domains")
	dbCommand.BoolVar(&args.Options.Sources, "src", false, "Print data sources for the discovered names")
	dbCommand.BoolVar(&args.Options.ASNTableSummary, "summary", false, "Print Just ASN Table Summary")
	dbCommand.BoolVar(&args.Options.DiscoveredNames, "names", false, "Print Just Discovered Names")
	dbCommand.BoolVar(&args.Options.NoColor, "nocolor", false, "Disable colorized output")
	dbCommand.BoolVar(&args.Options.ShowAll, "show", false, "Print the results for the enumeration index + domains provided")
	dbCommand.BoolVar(&args.Options.Silent, "silent", false, "Disable all output during execution")
	dbCommand.StringVar(&args.Filepaths.ConfigFile, "config", "", "Path to the INI configuration file. Additional details below")
	dbCommand.StringVar(&args.Filepaths.Directory, "dir", "", "Path to the directory containing the graph database")
	dbCommand.StringVar(&args.Filepaths.Domains, "df", "", "Path to a file providing root domain names")
	dbCommand.StringVar(&args.Filepaths.JSONOutput, "json", "", "Path to the JSON output file")
	dbCommand.StringVar(&args.Filepaths.TermOut, "o", "", "Path to the text file containing terminal stdout/stderr")

	if len(clArgs) < 1 {
		commandUsage(dbUsageMsg, dbCommand, dbBuf)
		return
	}

	if err := dbCommand.Parse(clArgs); err != nil {
		r.Fprintf(color.Error, "%v\n", err)
		os.Exit(1)
	}
	if help1 || help2 {
		commandUsage(dbUsageMsg, dbCommand, dbBuf)
		return
	}

	if args.Options.NoColor {
		color.NoColor = true
	}
	if args.Options.Silent {
		color.Output = ioutil.Discard
		color.Error = ioutil.Discard
	}

	if args.Filepaths.Domains != "" {
		list, err := config.GetListFromFile(args.Filepaths.Domains)
		if err != nil {
			r.Fprintf(color.Error, "Failed to parse the domain names file: %v\n", err)
			return
		}
		args.Domains.InsertMany(list...)
	}

	cfg := new(config.Config)
	cfg.LocalDatabase = true
	// Check if a configuration file was provided, and if so, load the settings
	if err := config.AcquireConfig(args.Filepaths.Directory, args.Filepaths.ConfigFile, cfg); err == nil {
		if args.Filepaths.Directory == "" {
			args.Filepaths.Directory = cfg.Dir
		}
		if len(args.Domains) == 0 {
			args.Domains.InsertMany(cfg.Domains()...)
		}
	} else if args.Filepaths.ConfigFile != "" {
		r.Fprintf(color.Error, "Failed to load the configuration file: %v\n", err)
		os.Exit(1)
	}

	db := openGraphDatabase(args.Filepaths.Directory, cfg)
	if db == nil {
		r.Fprintln(color.Error, "Failed to connect with the database")
		os.Exit(1)
	}
	defer db.Close()

	// Create the in-memory graph database for events that have information in scope
	memDB, err := memGraphForScope(args.Domains.Slice(), db)
	if err != nil {
		r.Fprintln(color.Error, err.Error())
		os.Exit(1)
	}

	// Get all the UUIDs for events that have information in scope
	uuids := memDB.EventList()
	if len(uuids) == 0 {
		r.Fprintln(color.Error, "Failed to find the domains of interest in the database")
		os.Exit(1)
	}

	if args.Options.ListEnumerations {
		listEvents(uuids, memDB)
		return
	}

	if args.Options.ShowAll || args.Filepaths.JSONOutput != "" {
		args.Options.DiscoveredNames = true
		args.Options.ASNTableSummary = true
	}

	if !args.Options.DiscoveredNames && !args.Options.ASNTableSummary {
		commandUsage(dbUsageMsg, dbCommand, dbBuf)
		return
	}

	// Put the events in chronological order
	uuids, _, _ = orderedEvents(uuids, memDB)
	if len(uuids) == 0 {
		r.Fprintln(color.Error, "Failed to sort the events")
		os.Exit(1)
	}

	// Select the enumeration that the user specified
	if args.Enum > 0 && len(uuids) >= args.Enum {
		idx := len(uuids) - args.Enum

		uuids = []string{uuids[idx]}
	}

	if args.Options.ASNTableSummary {
		fgY.Fprintln(color.Error, "Could take a moment while acquiring AS network information")
		// Migrate the changes back to the persistent db
		if healASInfo(uuids, memDB) {
			memDB.MigrateEvents(db, uuids...)
		}
	}

	showEventData(&args, uuids, memDB)
}

func listEvents(uuids []string, db *graph.Graph) {
	events, earliest, latest := orderedEvents(uuids, db)
	// Check if the user has requested the list of enumerations
	for pos, idx := 0, len(events)-1; idx >= 0; idx-- {
		if pos != 0 {
			g.Println()
		}

		g.Printf("%d) %s -> %s: ", pos+1, earliest[idx].Format(timeFormat), latest[idx].Format(timeFormat))
		// Print out the scope for this enumeration
		for x, domain := range db.EventDomains(events[idx]) {
			if x != 0 {
				g.Print(", ")
			}
			g.Print(domain)
		}
		g.Println()
		pos++
	}
}

func showEventData(args *dbArgs, uuids []string, db *graph.Graph) {
	var total int
	var err error
	var outfile *os.File
	var discovered []*requests.Output
	domains := args.Domains.Slice()

	if args.Filepaths.TermOut != "" {
		outfile, err = os.OpenFile(args.Filepaths.TermOut, os.O_WRONLY|os.O_CREATE, 0644)
		if err != nil {
			r.Fprintf(color.Error, "Failed to open the text output file: %v\n", err)
			os.Exit(1)
		}
		defer func() {
			outfile.Sync()
			outfile.Close()
		}()
		outfile.Truncate(0)
		outfile.Seek(0, 0)
	}

	tags := make(map[string]int)
	asns := make(map[int]*format.ASNSummaryData)
	for _, out := range getEventOutput(uuids, db) {
		if len(domains) > 0 && !domainNameInScope(out.Name, domains) {
			continue
		}

		out.Addresses = format.DesiredAddrTypes(out.Addresses, args.Options.IPv4, args.Options.IPv6)
		if l := len(out.Addresses); (args.Options.IPs || args.Options.IPv4 || args.Options.IPv6) && l == 0 {
			continue
		} else if l > 0 {
			total++
			format.UpdateSummaryData(out, tags, asns)
		}

		source, name, ips := format.OutputLineParts(out, args.Options.Sources,
			args.Options.IPs || args.Options.IPv4 || args.Options.IPv6, args.Options.DemoMode)
		if ips != "" {
			ips = " " + ips
		}

		if args.Options.DiscoveredNames {
			var written bool
			if outfile != nil {
				fmt.Fprintf(outfile, "%s%s%s\n", source, name, ips)
				written = true
			}
			if args.Filepaths.JSONOutput != "" {
				discovered = append(discovered, out)
				written = true
			}
			if !written {
				fmt.Fprintf(color.Output, "%s%s%s\n", blue(source), green(name), yellow(ips))
			}
		}
	}

	if total == 0 {
		r.Println("No names were discovered")
		return
	}
	if args.Filepaths.JSONOutput != "" {
		writeJSON(args, uuids, discovered, db)
	} else if args.Options.ASNTableSummary {
		var out io.Writer
		status := color.NoColor

		if outfile != nil {
			out = outfile
			color.NoColor = true
		} else if args.Options.ShowAll {
			out = color.Error
		} else {
			out = color.Output
		}

		format.FprintEnumerationSummary(out, total, tags, asns, args.Options.DemoMode)
		color.NoColor = status
	}
}

type jsonEvent struct {
	UUID   string `json:"uuid"`
	Start  string `json:"start"`
	Finish string `json:"finish"`
}

type jsonDomain struct {
	Domain string             `json:"domain"`
	Total  int                `json:"total"`
	Names  []*requests.Output `json:"names"`
}

type jsonOutput struct {
	Events  []*jsonEvent  `json:"events"`
	Domains []*jsonDomain `json:"domains"`
}

func writeJSON(args *dbArgs, uuids []string, assets []*requests.Output, db *graph.Graph) {
	var output jsonOutput

	// Add the event data to the JSON
	events, earliest, latest := orderedEvents(uuids, db)
	for i, uuid := range events {
		output.Events = append(output.Events, &jsonEvent{
			UUID:   uuid,
			Start:  earliest[i].Format(timeFormat),
			Finish: latest[i].Format(timeFormat),
		})
	}
	// Add the asset specific data
	for _, asset := range assets {
		var found bool
		var d *jsonDomain

		for _, domain := range output.Domains {
			if domain.Domain == asset.Domain {
				found = true
				d = domain
				break
			}
		}
		if !found {
			d = &jsonDomain{Domain: asset.Domain}

			output.Domains = append(output.Domains, d)
		}

		d.Total++
		d.Names = append(d.Names, asset)
	}

	jsonptr, err := os.OpenFile(args.Filepaths.JSONOutput, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		r.Fprintf(color.Error, "Failed to open the JSON output file: %v\n", err)
		return
	}
	// Remove previously stored data and encode the JSON
	jsonptr.Truncate(0)
	jsonptr.Seek(0, 0)
	json.NewEncoder(jsonptr).Encode(output)
	jsonptr.Sync()
	jsonptr.Close()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	_ "net/http/pprof"
	"os"

	"github.com/pajlada/pajbot2/pkg/common"
	"github.com/pajlada/pajbot2/pkg/common/config"
)

var buildTime string

var version = flag.Bool("version", false, "Show pajbot2 version")
var configPath = flag.String("config", "./config.json", "")

var validURLs = []string{
	"imgur.com",        // Image host
	"twitter.com",      // Social media
	"twimg.com",        // Twitter image host
	"forsen.tv",        // Bot website
	"pajlada.se",       // Bot creator website
	"pajlada.com",      // Bot creator website
	"pajbot.com",       // Bot website
	"youtube.com",      // Video hosting website
	"youtu.be",         // Youtube short-url
	"prntscr.com",      // Image host
	"prnt.sc",          // prntscr short-url
	"steampowered.com", // Game shop
	"gyazo.com",        // Image host
	"www.com",          // Meme
}

func main() {
	common.BuildTime = buildTime

	flag.Usage = func() {
		helpCmd()
	}
	flag.Parse()
	command := flag.Arg(0)

	if *version {
		fmt.Println(*version)
		os.Exit(0)
	}

	switch command {
	case "check":
		_, err := config.LoadConfig(*configPath)
		if err != nil {
			fmt.Println("An error occured while loading the config file:", err)
			os.Exit(1)
		} else {
			fmt.Println("No errors found in the config file")
			os.Exit(0)
		}

	case "install":
		installCmd()

	case "create":
		createCmd()

	case "help":
		helpCmd()

	default:
		fallthrough
	case "run":
		runCmd()
	}
}

func helpCmd() {
	_, err := os.Stderr.WriteString(
		`usage: pajbot2 <command> [<args>]
Commands:
   run            Run the bot (Default)
   check          Check the config file for missing fields
   install        Start the installation process (WIP)
   create <name>  Create a migration (WIP)
   newbot         Create a new bot
   linkchannel    Link a channel to a bot ID
`)
	if err != nil {
		log.Fatal(err)
	}
}

func runCmd() {
	application := NewApplication()

	err := application.LoadConfig(*configPath)
	if err != nil {
		log.Fatal("An error occured while loading the config file: ", err)
	}

	err = application.RunDatabaseMigrations()
	if err != nil {
		log.Fatal("An error occured while running database migrations: ", err)
	}

	err = application.StartRedisClient()
	if err != nil {
		log.Fatal("Error starting redis client:", err)
	}

	err = application.StartSQLClient()
	if err != nil {
		log.Fatal("Error starting SQL client:", err)
	}

	err = application.InitializeAPIs()
	if err != nil {
		log.Fatal("An error occured while initializing APIs: ", err)
	}

	err = application.LoadExternalEmotes()
	if err != nil {
		log.Fatal("An error occured while loading external emotes: ", err)
	}

	err = application.StartTwitterStream()
	if err != nil {
		log.Println("Error starting twitter stream: ", err)
	}

	err = application.StartWebServer()
	if err != nil {
		log.Fatal("An error occured while starting the web server: ", err)
	}

	/*
		err = application.LoadOldPajbot()
		if err != nil {
			log.Fatal("An error occured while loading old pajbot: ", err)
		}
	*/

	err = application.LoadBots()
	if err != nil {
		log.Fatal("An error occured while loading bots: ", err)
	}

	/*
		err = application.StartContextBot()
		if err != nil {
			log.Fatal("An error occured while starting context bot: ", err)
		}
	*/

	err = application.StartBots()
	if err != nil {
		log.Fatal("An error occured while starting bots: ", err)
	}

	err = application.StartPubSubClient()
	if err != nil {
		log.Fatal("An error occured while starting pubsub client", err)
	}

	log.Fatal(application.Run())
}

func installCmd() {
	_, err := os.Stderr.WriteString(
		`"install" not yet implemented
`)
	if err != nil {
		log.Fatal(err)
	}
}

func createCmd() {
	_, err := os.Stderr.WriteString(
		`"create" not yet implemented
`)
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>make pubsub client optional<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	_ "net/http/pprof"
	"os"

	"github.com/pajlada/pajbot2/pkg/common"
	"github.com/pajlada/pajbot2/pkg/common/config"
)

var buildTime string

var version = flag.Bool("version", false, "Show pajbot2 version")
var configPath = flag.String("config", "./config.json", "")

var validURLs = []string{
	"imgur.com",        // Image host
	"twitter.com",      // Social media
	"twimg.com",        // Twitter image host
	"forsen.tv",        // Bot website
	"pajlada.se",       // Bot creator website
	"pajlada.com",      // Bot creator website
	"pajbot.com",       // Bot website
	"youtube.com",      // Video hosting website
	"youtu.be",         // Youtube short-url
	"prntscr.com",      // Image host
	"prnt.sc",          // prntscr short-url
	"steampowered.com", // Game shop
	"gyazo.com",        // Image host
	"www.com",          // Meme
}

func main() {
	common.BuildTime = buildTime

	flag.Usage = func() {
		helpCmd()
	}
	flag.Parse()
	command := flag.Arg(0)

	if *version {
		fmt.Println(*version)
		os.Exit(0)
	}

	switch command {
	case "check":
		_, err := config.LoadConfig(*configPath)
		if err != nil {
			fmt.Println("An error occured while loading the config file:", err)
			os.Exit(1)
		} else {
			fmt.Println("No errors found in the config file")
			os.Exit(0)
		}

	case "install":
		installCmd()

	case "create":
		createCmd()

	case "help":
		helpCmd()

	default:
		fallthrough
	case "run":
		runCmd()
	}
}

func helpCmd() {
	_, err := os.Stderr.WriteString(
		`usage: pajbot2 <command> [<args>]
Commands:
   run            Run the bot (Default)
   check          Check the config file for missing fields
   install        Start the installation process (WIP)
   create <name>  Create a migration (WIP)
   newbot         Create a new bot
   linkchannel    Link a channel to a bot ID
`)
	if err != nil {
		log.Fatal(err)
	}
}

func runCmd() {
	application := NewApplication()

	err := application.LoadConfig(*configPath)
	if err != nil {
		log.Fatal("An error occured while loading the config file: ", err)
	}

	err = application.RunDatabaseMigrations()
	if err != nil {
		log.Fatal("An error occured while running database migrations: ", err)
	}

	err = application.StartRedisClient()
	if err != nil {
		log.Fatal("Error starting redis client:", err)
	}

	err = application.StartSQLClient()
	if err != nil {
		log.Fatal("Error starting SQL client:", err)
	}

	err = application.InitializeAPIs()
	if err != nil {
		log.Fatal("An error occured while initializing APIs: ", err)
	}

	err = application.LoadExternalEmotes()
	if err != nil {
		log.Fatal("An error occured while loading external emotes: ", err)
	}

	err = application.StartTwitterStream()
	if err != nil {
		log.Println("Error starting twitter stream: ", err)
	}

	err = application.StartWebServer()
	if err != nil {
		log.Fatal("An error occured while starting the web server: ", err)
	}

	/*
		err = application.LoadOldPajbot()
		if err != nil {
			log.Fatal("An error occured while loading old pajbot: ", err)
		}
	*/

	err = application.LoadBots()
	if err != nil {
		log.Fatal("An error occured while loading bots: ", err)
	}

	/*
		err = application.StartContextBot()
		if err != nil {
			log.Fatal("An error occured while starting context bot: ", err)
		}
	*/

	err = application.StartBots()
	if err != nil {
		log.Fatal("An error occured while starting bots: ", err)
	}

	err = application.StartPubSubClient()
	if err != nil {
		fmt.Println("Error starting PubSub Client:", err)
	}

	log.Fatal(application.Run())
}

func installCmd() {
	_, err := os.Stderr.WriteString(
		`"install" not yet implemented
`)
	if err != nil {
		log.Fatal(err)
	}
}

func createCmd() {
	_, err := os.Stderr.WriteString(
		`"create" not yet implemented
`)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>